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

# Simulate planet around stationary star 

using Plots # for plotting trajectory
using Measures # for adding margins to the plots (no cut-off labels)
using DifferentialEquations # for solving ODEs
using LinearAlgebra
using Unitful
using UnitfulAstro # astrophysical units

G = 4.0*pi^2 # time scale = year and length scale = AU
δ = 0.0 # deviation from inverse-square law


# GM_S = ustrip(uconvert(u"AU^3/yr^2", 1 * UnitfulAstro.GMsun)) # gravitational constant times mass of the Sun
GM_S = 4.0*pi^2 # gravitational constant times mass of the Sun in AU^3/yr^2 (same as above)
a = 0.39 # semi-major axis of Mercury
e = 0.206 # eccentricity of Mercury

t_final = 2 # final time of simulation in years

α_values = [.0004, .0008, .0016, .0032] # values of α to predict Mercury percession with

function tendency!(drp, rp, param, t)
   
   # 6d phase space
   r = rp[1:3] 
   p = rp[4:6] 

   (G, δ, α) = param

   r2 = dot(r, r)

   correction = (1 + α / r2) # correction to force law
   a = - G * r * r2^(-1.5-δ) * correction  # acceleration with M_S = 1

   drp[1:3] = p[1:3]
   drp[4:6] = a[1:3]
end

function energy(rp, param)

   r = rp[1:3]
   p = rp[4:6]

   (G, δ, α) = param

   r2 = dot(r, r)

   # TODO: add correction to potential energy
   pe = -G * r2^(-0.5 - δ)/(1.0 + 2.0 * δ)
   ke = 0.5 * dot(p, p)

   return pe + ke

end

function angularMomentum(rp)

   r = rp[1:3]
   p = rp[4:6]
   
   return cross(r, p)

end

function calculate_vy1_0(a, e, GM_S)
   return sqrt(GM_S * (1.0 - e)/(a * (1.0 + e)))
end

function calculate_x_0(a, e)
   return a * (1.0 + e)
end

function find_distance_derivative_changes(sol)
   i_change = []

   xs = sol[1, :]
   ys = sol[2, :]
   r = sqrt.(xs.^2 + ys.^2)
   derivatives = []
   for i in 1:length(r)-1
      push!(derivatives, r[i+1] - r[i])
   end

   for i in 1:length(derivatives)-1
      if derivatives[i] > 0 && derivatives[i+1] < 0 # if the derviative goes from positive to negative, store it
         if abs(derivatives[i]) < abs(derivatives[i+1])
            push!(i_change, i)
         else
            push!(i_change, i+1)
         end
      end
   end

   return i_change
end

function linear_regression(x, y)
    n = length(x)
    x̄ = sum(x) / n
    ȳ = sum(y) / n
    a = sum((x .- x̄) .* (y .- ȳ)) / sum((x .- x̄).^2)
    b = ȳ - a * x̄
    return (a, b)
end

function do_α_simulations(α_values, a, e, GM_S, t_final, δ)
   plots = []
   dθvdt = []

   x_0 = calculate_x_0(a, e)
   vy1_0 = calculate_vy1_0(a, e, GM_S)
   for α in α_values
      param = (GM_S, δ, α) # parameters of force law
      r0 = [x_0, 0.0, 0.0] # initial position in AU
      p0 = [0.0, vy1_0, 0.0] # initial velocity in AU / year
      rp0 = vcat(r0, p0) # initial condition in phase space 

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


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

      # get the times where the distance derivatives change
      i_changes = find_distance_derivative_changes(sol)

      # find the angle of percession at these derviate changes
      rad_to_deg = 180/π
      all_θ = [atan(sol[2, i] / sol[1, i]) * rad_to_deg for i in 1:length(sol.t)]

      # Plot of orbit
      xy = plot(
         sol[1, :], sol[2, :], 
         xlabel = "x (AU)", 
         ylabel = "y (AU)", 
         title = "Percession when α=$(α)", 
         aspect_ratio=:equal, 
         label="Orbit", 
         legend=:bottomright, 
         lw=.5,
         left_margin=5mm
         )
      # Add the position of the Sun
      scatter!(xy, [0], [0], label = "Sun")
      # Add the positions where the distance changes
      scatter!(
         xy, 
         sol[1, :][i_changes], 
         sol[2, :][i_changes], 
         label = "Dist. Derivative Changes", 
         color=:green
      )
      # Add to plots array
      push!(plots, xy)

      # Plot the change in angles over time
      dθ = scatter(
         sol.t[i_changes], 
         all_θ[i_changes], 
         xlabel = "time (yr)", 
         ylabel = "θ (degrees)", 
         title = "Orientation vs. Time (α=$(α))", 
         label="Dist. Derivative Changes", 
         color=:green,
         right_margin=5mm
      )
      # Add the linear regression of the change in angles
      (a, b) = linear_regression(sol.t[i_changes], all_θ[i_changes])
      a = round(a, digits=2)
      b = round(b, digits=2)
      plot!(dθ, sol.t[i_changes], a * sol.t[i_changes] .+ b, label = "θ ≈ $a * t + $b", left_margin=7mm)
      # Push to the plots array
      push!(plots, dθ)

      # Store the slope into the return array
      push!(dθvdt, a)
   end

   return plots, dθvdt
end

# Run the simulations
(plots, dθvdt) = do_α_simulations(α_values, a, e, GM_S, t_final, δ)

println("α_values = ", α_values)
println("dθvdt = ", dθvdt)

# Combine the plots into one plot
p_orbits = plot(plots..., layout=(4, 2), size=(800, 1000))
savefig(p_orbits, "hw4/4-10.png")

# Plot the change in dθ/dt over α
p_dθvdt = scatter(α_values, dθvdt, xlabel="α", ylabel="dθ/dt (degrees/year)", title="Percession Rate vs. α", lw=2, label="dθ/dt (from regression slopes)")
# Perform a linear regression on the data
(a, b) = linear_regression(α_values, dθvdt)
a = round(a, digits=2)
b = round(b, digits=2)
plot!(p_dθvdt, α_values, a * α_values .+ b, label = "dθ/dt ≈ $a * α + $b", left_margin=7mm)
savefig(p_dθvdt, "hw4/4-10-dθvdt.png")


percession_rate = a * 1.1e-8 # degrees per year to arcsec per year
# convert from degrees per year to arcsec per century
percession_rate = percession_rate * 3600 * 100 # arcsec per century
println("Percession rate = ", percession_rate, " arcsec/century")