#!/Applications/Julia-1.8.app/Contents/Resources/julia/bin/julia
# FOR PROBLEM 4.19
# author: sotech117

# Simulate planet around stationary star 

using Plots # for plotting trajectory
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

param = (G, δ) # parameters of force law

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


   ω = rp[8]

   (G, δ) = param

   r2 = dot(r, r)

   a = -G * r * r2^(-1.5-δ) # acceleration with m = 1

   x_c = r[1]
   y_c = r[2]

   τ = -3.0 * G * r2^(-2.5 - δ) * (x_c * sin(θ) - y_c * cos(θ)) * (x_c * cos(θ) + y_c * sin(θ))

   drp[1:3] = p[1:3]
   drp[4:6] = a[1:3]
   drp[7] = ω
   drp[8] = τ
end

function energy(rp, param)

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

   (G, δ) = param

   r2 = dot(r, r)

   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

# 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 simulate(r0x, p0y, θ01)
    r0 = [r0x, 0.0, 0.0] # initial position in HU 
    p0 = [0.0, p0y, 0.0] # initial velocity in HU / year
    θ0 = [θ01]
    pθ0 = [0.0]
    rp0 = vcat(r0, p0, θ0, pθ0) # initial condition in phase space 
    t_final = 10.0 # final time of simulation

    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

    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))
    println("Initial angular momentum = ", angularMomentum(sol[:,1]))
    println("Final angular momentum  = ", angularMomentum(sol[:, end]))

    # Plot θ over time
    sol[7, :] = clean_θ(sol[7, :])

    return sol
end

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

function calculate_Δθ(sol1, sol2, ticks=1000)
    # generate a equal set of times from the smaller of the two simulations
    time_0 = min(sol1.t[1], sol2.t[1])
    time_f = min(sol1.t[end], sol2.t[end])
    # make an even step fucntion over the length of these times
    sample_times = range(time_0, time_f, length=ticks)

    # interpolate the two simulations to the same time
    diff = []
    for i in 1:length(sample_times)-1
        t = sample_times[i]

        # find the index that is before the time
        idx1 = -1
        for j in 1:length(sol1.t)
            if sol1.t[j] > t
                idx1 = j - 1
                break
            end
        end
        idx2 = -1
        for j in 1:length(sol2.t)
            if sol2.t[j] > t
                idx2 = j - 1
                break
            end
        end

        # interpolate the values between the two timestamps
        θ1 = sol1[7, idx1] + (sol1[7, idx1+1] - sol1[7, idx1]) * (t - sol1.t[idx1])/(sol1.t[idx1+1] - sol1.t[idx1])
        θ2 = sol2[7, idx2] + (sol2[7, idx2+1] - sol2[7, idx2]) * (t - sol2.t[idx2])/(sol2.t[idx2+1] - sol2.t[idx2])

        push!(diff, sqrt((θ1 - θ2)^2))
    end

    return (sample_times[1:end-1], diff)
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 find_l_exponent(eccentricity, GM_S)
    a = 1.0
    vy1_0 = calculate_vy1_0(a, eccentricity, GM_S)
    sol1 = simulate(1.0, vy1_0, 0.0)
    sol2 = simulate(1.0, vy1_0, 0.01) # small change in θ_0
    (ts, diff_θ) = calculate_Δθ(sol1, sol2)

    plt = plot(ts, diff_θ, yscale=:ln, xlabel="time (s)", ylabel="Δθ (radians)", title="Δθ vs time", legend=:bottomright, lw=2, label="\$ Δ θ ≡ √{(θ_1 - θ_2)^2} \$")

    # linear regression of the log of the difference
    (a, b) = linear_regression(ts, log.(diff_θ))
    a_rounded = round(a, digits=4)
    b = round(b, digits=4)
    plt = plot!(ts, exp.(a*ts .+ b), label="exp($a_rounded*t + $b)", lw=1.2, color=:red, linestyle=:dash)

    return a, plt
end

function process_data!(les, es, plt)
    # find the first index where the l exponent is greater than .2
    idx = -1
    for i in 1:length(les)
        if les[i] > .2
            idx = i
            break
        end
    end

    println("e=$(es[idx]) & l=$(les[idx])")

    # add a vertical line at this value of e
    e_rounded = round(es[idx], digits=2)
    plot!(plt, [es[idx], es[idx]], [-0, .75], label="e : L > 0 ≈ $(e_rounded)", color=:red, lw=.75, linestyle=:dash)

    # perform a log regression only on values where the l exponent is greater than .2
    # (a, b) = linear_regression(es[idx:end], log.(les[idx:end]))
    # a_rounded = round(a, digits=2)
    # b = round(b, digits=2)
    # plot!(plt, es[idx:end], exp.(a*es[idx:end] .+ b), label="L ≈ exp($a_rounded*e + $b)", lw=2, color=:orange, linestyle=:dash)

    return
end

# simualte two elipical
# eccentricity = .224
# (a, plt) = find_l_exponent(eccentricity, G)

# generate a span of eccentricities

# eccentricities = range(0.0, .985, length=750) # todo: add more len for final render
# l_exponents = []
# for e in eccentricities
#     (a, _) = find_l_exponent(e, G)
#     push!(l_exponents, a)
# end

# # do a normal plot
# # p1 = plot(
# #     eccentricities, l_exponents, xlabel="eccentricity (e)", ylabel="Lyapunov Exponents (L)", title="Lyapunov Exponent vs Eccentricity", 
# #     lw=.8, label="Lyapunov Exponents (connected)")
# p1 = scatter(eccentricities, l_exponents, 
# title="Lyapunov Exponent vs Eccentricity", xlabel="Eccentricity (e)", ylabel="Lyapunov Exponent (L)", 
# color=:blue, markerstrokealpha=0.0, markersize=2, msw=0.0,
# label="Lyapunov Exponents")

# # process the data
# process_data!(l_exponents, eccentricities, p1)

# savefig(p1, "hw4/4-19-final.png")


eccentricity = .224
vy_elipitical = calculate_vy1_0(1.0, eccentricity, G)
println("vy1_0 = ", vy_elipitical)
sol1 = simulate(1.0, vy_elipitical, 0.0)
sol2 = simulate(1.0, vy_elipitical, 0.01) # small change in θ_0
(ts, diff_θ) = calculate_Δθ(sol1, sol2)
# plot the θ over time
plt_1 = plot(sol1.t, sol1[7, :], xlabel="time (Hyperion years)", ylabel="θ (radians)", title="θ vs time \$(θ_0 = 0, e = .224\$)", legend=false, lw=2, label="θ_0 = 0.0")
# plot the ω over time
plt_2 = plot(sol2.t, sol2[7, :], xlabel="time (Hyperion years)", ylabel="θ (radians)", title="θ vs time \$(θ_0 = .01, e = .224)\$", legend=false, lw=2, label="ω_0 = 0.0")
# plot the difference in θ over time
plot_3 = plot(ts, diff_θ, yscale=:ln, xlabel="time (Hyperion years) ", ylabel="Δθ (radians)", title="Δθ vs time", legend=:bottomright, lw=2, label="\$ Δ θ ≡ √{(θ_1 - θ_2)^2} \$")

# apply a linear regression to the log of the difference
(a, b) = linear_regression(ts, log.(diff_θ))
a_rounded = round(a, digits=2)
b = round(b, digits=2)
plot_3 = plot!(ts, exp.(a*ts .+ b), label="exp($a_rounded*t + $b)", lw=1.2, color=:red, linestyle=:dash)

plt_combined = plot(plt_1, plt_2, plot_3, layout=(3, 1), size=(800, 1000))
savefig(plt_combined, "hw4/4-19-.224.png")

# run another simuation, but only a circle
e = 0
vy_circular = calculate_vy1_0(1.0, e, G)
println("vy1_0 = ", vy_circular)
sol1 = simulate(1.0, vy_circular, 0.0)
sol2 = simulate(1.0, vy_circular, 0.01) # small change in θ_0
(ts, diff_θ) = calculate_Δθ(sol1, sol2)
# plot the θ over time
plt_1 = plot(sol1.t, sol1[7, :], xlabel="time (Hyperion years)", ylabel="θ (radians)", title="θ vs time (\$θ_0 = 0, e = 0\$)", lw=2, legend=false)
# plot the ω over time
plt_2 = plot(sol2.t, sol2[7, :], xlabel="time (Hyperion years)", ylabel="θ (radians)", title="θ vs time (\$θ_0 = .01, e = 0)\$", lw=2, legend = false)
# plot the difference in θ over time
plt = plot(ts, diff_θ, yscale=:ln, xlabel="time (Hyperion years)", ylabel="Δθ (radians)", title="Δθ vs time", legend=:bottomright, lw=2, label="\$ Δ θ ≡ √{(θ_1 - θ_2)^2} \$")

# apply a linear regression to the log of the difference
(a, b) = linear_regression(ts, log.(diff_θ))
a_rounded = round(a, digits=2)
b = round(b, digits=2)
plt = plot!(ts, exp.(a*ts .+ b), label="exp($a_rounded*t + $b)", lw=1.2, color=:red, linestyle=:dash)

plt_combined = plot(plt_1, plt_2, plt, layout=(3, 1), size=(800, 1000))

savefig(plt_combined, "hw4/4-19-0.png")