# FOR PHYS2600 HW3
# author: sotech117

using DifferentialEquations
using Plots

g = 9.8
t_span = (0.0, 1000.0)
p = g

function tendency!(du, u, p, t)
    q_1, q_2, p_1, p_2 = u
    g = p

    denominator = 1.0 + (sin(q_1 - q_2)*sin(q_1 - q_2))

    dq_1 = (p_1 - p_2 * cos(q_1 - q_2)) / denominator
    dq_2 = (-p_1*cos(q_1 - q_2) + 2*p_2) / denominator

    k_1 = p_1*p_2*sin(q_1 - q_2) / denominator
    k_2 = (p_1*p_1 + 2*p_2*p_2 - 2*p_1*p_2*cos(q_1 - q_2)) / (2*denominator*denominator)

    dp_1 = -2 * sin(q_1) - k_1 + k_2 * sin(2*(q_1 - q_2))
    dp_2 = - sin(q_2) + k_1 - k_2 * sin(2*(q_1 - q_2))
    
    du[1] = dq_1
    du[2] = dq_2
    du[3] = dp_1
    du[4] = dp_2
end

function energy(u, p)
    q_1, q_2, p_1, p_2 = u
    g = p

    T = 0.5 * (p_1*p_1 + p_2*p_2)
    V = - 2 * cos(q_1) - cos(q_2)
    return T + V
end

function build_ensemble(q_start, q_end, p_start, p_end, n)
    q_1s = zeros(n)
    q_2s = zeros(n)
    p_1s = zeros(n)
    p_2s = zeros(n)
    for i in 1:n
        q_1s[i] = q_start + (q_end - q_start) * rand()
        q_2s[i] = q_start + (q_end - q_start) * rand()
        p_1s[i] = p_start + (p_end - p_start) * rand()
        p_2s[i] = p_start + (p_end - p_start) * rand()
    end
    return [q_1s, q_2s, p_1s, p_2s]
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 run_simulation(u_0, tspan, p)
    prob = ODEProblem(tendency!, u_0, tspan, p)
    sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8) # control simulation
    sol[1, :] = clean_θ(sol[1, :])
    sol[2, :] = clean_θ(sol[2, :])
    return sol
end

# returns the new ensemble after t_span time
function simulate_ensembles(ensemble, tspan, p)
    q_1s = zeros(length(ensemble[1]))
    q_2s = zeros(length(ensemble[2]))
    p_1s = zeros(length(ensemble[3]))
    p_2s = zeros(length(ensemble[4]))
    println("length ensemble = ", length(ensemble[1]))
    for i in 1:length(ensemble[1])
        u_0 = [ensemble[1][i], ensemble[2][i], ensemble[3][i], ensemble[4][i]]
        sol = run_simulation(u_0, tspan, p)

        q_1s[i] = sol[1, end]
        q_2s[i] = sol[2, end]
        p_1s[i] = sol[3, end]
        p_2s[i] = sol[4, end]
    end
    println("q1_s", length(q_1s))
    return [q_1s, q_2s, p_1s, p_2s]
end

function sum_all_positions_and_momenta(ensemble)
    q_1 = 0
    q_2 = 0
    p_1 = 0
    p_2 = 0
    for i in 1:length(ensemble[1])
        q_1 += ensemble[1][i]
        q_2 += ensemble[2][i]
        p_1 += ensemble[3][i]
        p_2 += ensemble[4][i]
    end
    return [q_1, q_2, p_1, p_2]
end

function estimate_phase_space_volume(ensemble, n)
    # make 100000 random points in the 4d phase space from [-pi, pi] x [-pi, pi] x [-3, 3] x [-3, 3]
    points = build_ensemble(-pi, pi, -3, 3, 100000)

    # iterate over the points and find the number of points that hit the ensemble
    count = 0
    for i in 1:length(points[1])
        for j in 1:length(ensemble[1])
            if abs(points[1][i] - ensemble[1][j]) < n && abs(points[2][i] - ensemble[2][j]) < n && abs(points[3][i] - ensemble[3][j]) < n && abs(points[4][i] - ensemble[4][j]) < n
                count += 1
                break
            end
        end
    end

    # return the volume
    return (count / 100000) # note units don't matter, just comparing
end


# PROBLEM 1 -> show that two systems with the same intial energy can have different chaotic results
u_1_0 = [0, 0, 6.26, 0.0]
u_2_0 = [0.0, 0.0, 0.0, 6.26]

# print the initial energies
println("energy 1 = ", energy(u_1_0, p))
println("energy 2 = ", energy(u_2_0, p))

sol_1 = run_simulation(u_1_0, t_span, p)
sol_2 = run_simulation(u_2_0, t_span, p)

# plot the phase space
p1 = scatter(sol_1[1, :], sol_1[3, :], label="q_1", xlabel="q_1 (radians)", ylabel="p_1 (momentum)", title="Phase Space i=1 for Pendulum 1", legend=false, color=:blue, msw=0, ms=.75)
p2 = scatter(sol_2[1, :], sol_2[3, :], label="q_2", xlabel="q_1 (radians)", ylabel="p_1 (momentum)", title="Phase Space i=1 for Pendulum 2", legend=false, color=:red, msw=0, ms=.75)

# plot the other phase space
p3 = scatter(sol_1[2, :], sol_1[4, :], label="q_1", xlabel="q_2 (radians)", ylabel="p_2 (momentum)", title="Phase Space i=2 for Pendulum 1", legend=false, color=:blue, msw=0, ms=.75)
p4 = scatter(sol_2[2, :], sol_2[4, :], label="q_2", xlabel="q_2 (radians)", ylabel="p_2 (momentum)", title="Phase Space i=2 for Pendulum 2", legend=false, color=:red, msw=0, ms=.75)

# plot the trajectories
p5 = scatter(sol_1[1, :], sol_1[2, :], label="q_1 (radians)", xlabel="q_1 (radians)", ylabel="q_2 (radians)", title="Relative Trajectory for Pend1", legend=false, color=:blue, msw=0, ms=.75)
p6 = scatter(sol_2[1, :], sol_2[2, :], label="q_2 (radians)", xlabel="q_1 (radians)", ylabel="q_2 (radians)", title="Relative Trajectory for Pend2", legend=false, color=:red, msw=0, ms=.75)

# plot the trajectories over time
p7 = plot(sol_1.t, sol_1[1, :], label="pendulum 1", xlabel="time (s)", ylabel="q_1 (radians)", title="q_1 vs time", legend=:topright)
plot!(p7, sol_2.t, sol_2[1, :], label="pendulum 2", xlabel="time (s)", ylabel="q_1 (radians)", title="q_1 vs time")

p8 = plot(sol_1.t, sol_1[2, :], label="pendulum 1", xlabel="time (s)", ylabel="q_2 (radians)", title="q_2 vs time", legend=:topright)
plot!(p8, sol_2.t, sol_2[2, :], label="pendulum 2", xlabel="time (s)", ylabel="q_2 (radians)", title="q_2 vs time")

plt = plot(p1, p2, p3, p4, p5, p6, p7, p8, layout=(4, 2), size=(1000, 1000))

savefig(plt, "hw3/double_pendulum.png")

# PROBLEM 2 -> show that the volume of the phase space is conserved

ensemble = build_ensemble(-1, 1, -1, 1, 1000)
# plot these points
s1 = scatter(ensemble[1], ensemble[3], label="t=0", xlabel="q_1 (radians)", ylabel="p_1 (momentum)", title="Phase Space (i=1)", color=:blue, msw=.5)
s2 = scatter(ensemble[2], ensemble[4], label="t=0", xlabel="q_2 (radians)", ylabel="p_2 (momentum)", title="Phase Space (i=2)", color=:red, msw=.5)
# simulate the ensemble
new_ensemble = simulate_ensembles(ensemble, t_span, p)
# plot these points
scatter!(s1, new_ensemble[1], new_ensemble[3], label="t=1000", xlabel="q_1 (radians)", ylabel="p_1 (momentum)", title="Phase Space (i=1)", color=:green, msw=.5)
scatter!(s2, new_ensemble[2], new_ensemble[4], label="t=1000", xlabel="q_2 (radians)", ylabel="p_2 (momentum)", title="Phase Space (i=2)", color=:green, msw=.5)

# print the volumes
println("\n\nOriginal Volume:\t", estimate_phase_space_volume(ensemble, .1))
println("Final Volume:\t", estimate_phase_space_volume(new_ensemble, .1))

plt_2 = plot(s1, s2)
savefig(plt_2, "hw3/ensemble.png")