using DifferentialEquations
using Plots


# Many-body, 2D.

const F = Float64
const D = 2
const G = 6.6743e-11  # m³/kg s²

const M_S = 2e30
const M_E = 6e24
const M_M = 7.35e22

const R_E = 150e9
const R_M = 400e6

mutable struct MBSystem
    state::Matrix{F}
    masses::Vector{F}
    G::F
end

# TODO: Write a recipe for this.
function plot_mbs(mbs::MBSystem)
    scatter(mbs.state[1,:]', mbs.state[2,:]', markersize=3 .+ 10 * log.(mbs.masses'), legend=:false)
    for i in 1:length(mbs.masses)
        p0 = mbs.state[1:D, i]
        p1 = p0 + mbs.state[D+1:end, i]

        plot!([p0[1], p1[1]], [p0[2], p1[2]], width=2)
    end
    current()
end

function plot_solution(mbs::MBSystem, solution)
    p0 = mbs.state[1:D, :]
    vars = [(2D*(i-1) + 1, 2D*(i-1) + 2) for i in 1:length(mbs.masses)]
    plot(solution, vars=vars, labels=[string(i) for j in [1], i in (1:length(mbs.masses))], size=(1000,1000))
    current()
end

function animate_solution(mbs::MBSystem, solution; filename="orbit.gif")
    p0 = mbs.state[1:D, :]
    vars = [(2D*(i-1) + 1, 2D*(i-1) + 2) for i in 1:length(mbs.masses)]
    anim = @animate for t in solution.t
        plot(solution, vars=vars, tspan=(0, t), labels=[string(i) for j in [1], i in (1:length(mbs.masses))], size=(1000,1000), width=5)
    end
    gif(anim, filename)
end

function new_mb(n::Int; masses=nothing, G=G)::MBSystem
    if !isnothing(masses) && length(masses) != n
        error("Length of masses must equal n")
    end
    MBSystem(zeros(2D, n), isnothing(masses) ? zeros(n) : masses, G)
end

function mb_set(mbs::MBSystem, i::Int, pos::Vector, spd::Vector=[0,0])
    mbs.state[1:D, i] .= pos;
    mbs.state[D+1:end, i] .= spd;
    nothing
end

function radial_equilibrium_speed(mbs::MBSystem, i::Int, r::F)
    # v²/r = G M / r² → v = √(G M / r)
    return sqrt(mbs.G * mbs.masses[i] / r) 
end

# Return ̂r / r^2 for bodies i, j where r = r_i - r_j
function radial_term(m::Matrix{F}, i::Int, j::Int)::Vector
    r = m[1:D, j] - m[1:D, i]
    l = sqrt(sum(r.^2))
    r ./ l^3
end

# Calculate vector of acceleration according to Newton's gravity law.
# Considers all bodies j ≠ i.
# m a = G (m M / r^2) * ̂r
function acceleration(mbs::MBSystem, i::Int; state=nothing)::Vector{F}
    n = length(mbs.masses)
    a = zeros(D)
    st = isnothing(state) ? mbs.state : state
    for j = 1:n
        if i == j
            continue
        end
        rt = radial_term(st, i, j)
        a .+= rt * mbs.G * mbs.masses[j]
    end
    a
end

function DE(du, u, p, t)
    # Link velocities. d{x,y}/dt = v{x,y}
    du[1:D, :] .= u[D+1:end, :]
    # Calculate force application
    for i in 1:length(p.masses)
        du[D+1:end, i] .= acceleration(p, i; state=u)
    end
end

function solve_sme()
    MB_SME = new_mb(3, masses=[M_S, M_E, M_M], G=G)
    mb_set(MB_SME, 1, [0., 0], [0., 0.])
    mb_set(MB_SME, 2, [0., -R_E], [1.1*radial_equilibrium_speed(MB_SME, 1, R_E), 0])
    mb_set(MBSME, 3, [0., -(R_E+R_M)], [1.1*radial_equilibrium_speed(MB_SME, 1, R_E) + radial_equilibrium_speed(MB0, 2, R_M), 0])

    problem = ODEProblem(DE, MB_SME.state, (0, 365*24*3600), MB_SME)
    solution = solve(problem)

    plot_solution(MB_SME, solution)
end

function solve_me()
    MB_ME = new_mb(2, masses=[M_E, M_M], G=G)
    mb_set(MB_ME, 1, [0., 0], [0., 0])
    mb_set(MB_ME, 2, [0., -R_M], [1.1*radial_equilibrium_speed(MB_ME, 1, R_M), 0])

    problem = ODEProblem(DE, MB_ME.state, (0, 365*24*3600), MB_ME)
    solution = solve(problem)

    plot_solution(MB_ME, solution)
end

function solve_3body()
    MB0 = new_mb(3, masses=[40., 2, 1], G=0.002)
    mb_set(MB0, 1, [0., 0.], [0., 0.])
    mb_set(MB0, 2, [0., 5.], [-0.15, 0.0])
    mb_set(MB0, 3, [0., -5.], [0.1, 0.0])

    problem = ODEProblem(DE, MB0.state, (0, 1000), MB0)
    solution = solve(problem)

    animate_solution(MB0, solution)
end