# L-30 MCS 471 Mon 31 Oct 2022 : rk4celestial.jl

using Printf

"""
    rk4(F::Function,
        h::Float64,endT::Float64,
        init::Array{Float64})

applies a 4-stage Runge-Kutta method with step size h,
for the time interval [0, endT] with initial values in init.

The function F has two input arguments:
1. a number, the value of the independent variable (time),
2. a vector with the values of the dependent variables.

Returns an array of size length(init) + 1 with arrays all
of the same length: the number of time steps.
"""
function rk4(F::Function,
             h::Float64,endT::Float64,
             init::Array{Float64})
    dim = length(init)
    k1 = zeros(dim)
    k2 = zeros(dim)
    k3 = zeros(dim)
    k4 = zeros(dim)
    n = Int(ceil(endT/h))
    result = [zeros(n+1) for k=1:dim+1]
    for k=1:dim
        result[k+1][1] = init[k]
    end
    y = init
    for i=1:n
        t = i*h
        k1 = F(t, y)
        k2 = F(t + h/2, y + (h/2)*k1)
        k3 = F(t + h/2, y + (h/2)*k2)
        k4 = F(t + h, y + h*k3)
        y = y + (h/6)*(k1 + 2*k2 + 2*k3 + k4)
        result[1][i+1] = t
        for k=1:dim
            result[k+1][i+1] = y[k]
        end
    end
    return result
end

"""
    Runs a test of rk4 on the pendulum.
"""
function testRK4pend()
    println("running RK4 on the pendulum ...")
    f(t, y) = [y[2], -9.807*sin(y[1])]
    T = 2*pi
    h = T/24
    r = rk4(f, h, T, [pi/4, 0.0])
    t, y1, y2 = r
    for i=1:length(t)
        stri = @sprintf("%3d", i)
        strt = @sprintf("%.2f", t[i])
        stry1 = @sprintf("%.6e", y1[i])
        stry2 = @sprintf("%.6e", y2[i])
        println("$stri  $strt  $stry1  $stry2")
    end
end

# testRK4pend()

"""
    Right hand size for the 3-body problem.
"""
function F(t, y)
    x1, u1, y1, v1, x2, u2, y2, v2, x3, u3, y3, v3 = y
    r = [u1, -(x1-x2)/((x1-x2)^2 + (y1-y2)^2)^(3/2) - (x1-x3)/((x1-x3)^2 + (y1-y3)^2)^(3/2),
         v1, -(y1-y2)/((x1-x2)^2 + (y1-y2)^2)^(3/2) - (y1-y3)/((x1-x3)^2 + (y1-y3)^2)^(3/2),
         u2, -(x2-x1)/((x2-x1)^2 + (y2-y1)^2)^(3/2) - (x2-x3)/((x2-x3)^2 + (y2-y3)^2)^(3/2),
         v2, -(y2-y1)/((x2-x1)^2 + (y2-y1)^2)^(3/2) - (y2-y3)/((x2-x3)^2 + (y2-y3)^2)^(3/2),
         u3, -(x3-x1)/((x3-x1)^2 + (y3-y1)^2)^(3/2) - (x3-x2)/((x3-x2)^2 + (y3-y2)^2)^(3/2),
         v3, -(y3-y1)/((x3-x1)^2 + (y3-y1)^2)^(3/2) - (y3-y2)/((x3-x2)^2 + (y3-y2)^2)^(3/2)]
    return r
end

"""
    Runs the 3-body problem ...
"""
function main()
    # the initial positions
    ip1 = [0.97000436, -0.24308753]
    ip2 = [-ip1[1], -ip1[2]] 
    ip3 = [0, 0]
    # the initial velocities
    iv3 = [-0.93240737, -0.86473146]
    iv2 = [-iv3[1]/2, -iv3[2]/2] 
    iv1 = iv2
    # The initial conditions
    init = [ip1[1], iv1[1], ip1[2], iv1[2], 
            ip2[1], iv2[1], ip2[2], iv2[2], 
            ip3[1], iv3[1], ip3[2], iv3[2]]
    y = F(0.0, init)
    println(y)
    # T = 6.326
    T = 3.142/2
    h = 0.01
    r = rk4(F, h, T, init)
    t, x1, u1, y1, v1, x2, u2, y2, v2, x3, u3, y3, v3 = r
    println("initial positions :")
    println(ip1[1], "  ", ip1[1])
    println(ip2[1], "  ", ip2[1])
    println(ip3[1], "  ", ip3[1])
    println("end positions at t = ", t[end])
    println(x1[end], "  ", y1[end])
    println(x2[end], "  ", y2[end])
    println(x3[end], "  ", y3[end])
end

main()