# FOR PROBLEM 3.12
# author: sotech117

#!/Applications/Julia-1.8.app/Contents/Resources/julia/bin/julia

# Simulate driven pendulum to find chaotic regime

using Plots # for plotting trajectory
using DifferentialEquations # for solving ODEs

ω0 = 1.0 # ω0^2 = g/l
β = 0.5 # β = friction
f = 1.2 # forcing amplitude
ω = .66667 # forcing frequency
param = (ω0, β, f, ω) # parameters of anharmonic oscillator

function tendency!(dθp::Vector{Float64}, θp::Vector{Float64}, param, t::Float64)
   
   (θ, p) = θp # 2d phase space
   (dθ, dp) = dθp # 2d phase space derviatives

   (ω0, β, f, ω) = param

   a = -ω0^2 * sin(θ) - β * dθ + f * forcing(t, ω) # acceleration with m = 1

   dθp[1] = p
   dθp[2] = a
end

function forcing(t::Float64, ω::Float64)

   return sin(ω * t)

end

function energy(θp::Vector{Float64}, param)

   (θ, p) = θp

   (ω0, β, f, ω) = param

   pe = ω0^2 * (1.0 - cos(θ))
   ke = 0.5 * p^2

   return pe + ke

end

# take a list and reduce theta to the interval [-π, π]
function clean_θ(θ::Vector{Float64})
   rθ = []
   for i in 1:length(θ)
      tmp = θ[i] % (2 * π)
      if tmp > π
         tmp = tmp - 2 * π
      elseif tmp < -π
         tmp = tmp + 2 * π
      end
      push!(rθ, tmp)
   end
   return rθ
end

function get_poincare_sections(sample_θ, sample_p, sample_t, Ω_d, ϵ::Float64, phase_shift=0.0::Float64)
   n = 0

   poincare_θ = []
   poincare_p = []

   for i in 1:length(sample_θ)
      if abs(sample_t[i] * Ω_d - (2 * π * n + phase_shift)) < ϵ / 2
         push!(poincare_θ, sample_θ[i])
         push!(poincare_p, sample_p[i])
         n += 1
      end
   end
   
   return (poincare_θ, poincare_p)
end

θ0 = 0.2 # initial position in meters
p0 = 0.0 # initial velocity in m/s
θp0 = [θ0, p0] # initial condition in phase space 
t_final = 1000.0 # final time of simulation

tspan = (0.0, t_final) # span of time to simulate

prob = ODEProblem(tendency!, θp0, tspan, param) # specify ODE
sol = solve(prob, Tsit5(), reltol=1e-12, abstol=1e-12) # solve using Tsit5 algorithm to specified accuracy

sample_times = sol.t
println("\n\t Results")
println("final time  = ", sample_times[end])
println("Initial energy = ", energy(sol[:,1], param))
println("Final energy = ", energy(sol[:, end], param))

(ω0, β, f, ω) = param

# Plot of position vs. time
# θt = plot(sample_times, [sol[1, :], f * forcing.(sample_times, ω)], xlabel = "t", ylabel = "θ(t)", legend = false, title = "θ vs. t")

# Phase space plot
cleaned = clean_θ(sol[1, :])
θp = scatter(cleaned, sol[2, :], xlabel = "θ (radians)", ylabel = "ω (radians/s)", legend = false, title = "Phase Space Plot", mc=:black, ms=.35, ma=1)


# plot the poincare sections
(poincare_θ, pointcare_p) = get_poincare_sections(cleaned, sol[2, :], sol.t, ω, 0.1)
s1 = scatter(poincare_θ, pointcare_p, xlabel = "θ (radians)", ylabel = "ω (radians/s)", label="2nπ", title = "Poincare Sections", mc=:red, ms=2, ma=0.75)
s2 = scatter(get_poincare_sections(cleaned, sol[2, :], sol.t, ω, 0.1, π / 2.0), mc=:blue, ms=2, ma=0.75, label="2nπ + π/2", title="Poincare Sections", xlabel = "θ (radians)", ylabel = "ω (radians/s)", legend=:bottomleft)
s3 = scatter(get_poincare_sections(cleaned, sol[2, :], sol.t, ω, 0.1, π / 4.0), mc=:green, ms=2, ma=0.75, label="2nπ + π/4", title="Poincare Sections", xlabel = "θ (radians)", ylabel = "ω (radians/s)")

plot(θp, s1, s2, s3)