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

# Simulate a solar system

using Plots # for plotting trajectory
using DifferentialEquations # for solving ODEs
using LinearAlgebra # for dot and cross products

G = 4.0*pi^2 # time scale = year and length scale = AU

mutable struct body
   name::String # name of star or planet
   m::Float64 # mass
   r::Vector{Float64} # position vector
   p::Vector{Float64} # momentum vector
end 

function angularMomentum(b::body)
   r = b.r
   p = b.p
   return cross(r, p)
end

function kineticEnergy(b::body)
   p = b.p
   m = b.m
   return dot(p, p) / (2.0 * m)
end

function rp(b::body)
   return [b.r; b.p]
end

function force(body1::body, body2::body)
   r = body1.r - body2.r
   rSquared = dot(r, r)
   return -G * body1.m * body2.m * r * rSquared^(-1.5)
end

function potentialEnergy(body1::body, body2::body)
   r = body1.r - body2.r
   rSquared = dot(r, r)
   return -G * body1.m * body2.m * rSquared^(-0.5)
end

mutable struct SolarSystem
   bodies::Vector{body}
   numberOfBodies::Int64
   phaseSpace::Matrix{Float64} # 6N dimensional phase space
end

function SolarSystem()

   bodies = Vector{body}()
   push!(bodies, body("Sun", 1.0, zeros(3), zeros(3)))
   #push!(bodies, body("Earth", 1.0, [-1.0, 0.0, 0.0], [0.0, 1.0 * pi, 0.0]))
   push!(bodies, body("Star", 1.0, [1.0, 0.0, 0.0], [0.0, sqrt(1.95) * 2.0 * pi, 0.0]))
   #push!(bodies, body("Jupiter", 1.0, [3.0, 0.0, 0.0], [0.0, 0.25 * pi, 0.0]))
   numberOfBodies = size(bodies)[1]

   phaseSpace = zeros(6, 0)
   for b in bodies
      phaseSpace = [phaseSpace rp(b)]
   end
   
   return SolarSystem(bodies, numberOfBodies, phaseSpace)

end

function TotalAngularMomentum(s::SolarSystem)
   L = zeros(3)
   for b in s.bodies
      L += angularMomentum(b)
   end

   return L
end

function TotalEnergy(s::SolarSystem)
   ke = 0.0
   pe = 0.0

   for body1 in s.bodies
      ke += kineticEnergy(body1)
      for body2 in s.bodies
         if (body1 != body2)
            pe += 0.5 * potentialEnergy(body1, body2)
         end
      end
   end

   return pe + ke
end

function ZeroOutLinearMomentum!(s::SolarSystem)

   totalLinearMomentum = zeros(3)
   totalMass = 0.0
   for body in s.bodies
      totalLinearMomentum += body.p
      totalMass += body.m
   end
   
   s.phaseSpace = zeros(6, 0)
   for body in s.bodies
      body.p -= body.m * totalLinearMomentum / totalMass
      s.phaseSpace = [s.phaseSpace rp(body)]
   end

   return nothing
end

function tendency!(dps, ps, s::SolarSystem, t)

   i = 1 # update phase space of individual bodies
   for b in s.bodies
      b.r = ps[1:3, i]
      b.p = ps[4:6, i]
      i += 1
   end

   # find velocities of bodies and forces on them. O(N^2) computational cost
   N = s.numberOfBodies
   for i in 1:N
      b1 = s.bodies[i]
      dps[1:3, i] = b1.p / b1.m  #dr/dt
      dps[4:6, i] = zeros(3)
      for j in 1:i-1
            b2 = s.bodies[j]
            f = force(b1, b2) # call only once
            dps[4:6, i] += f
            dps[4:6, j] -= f # Newton's 3rd law
      end
   end

   return nothing
end

s = SolarSystem()
ZeroOutLinearMomentum!(s)
println(typeof(s))
println("Initial total energy = ", TotalEnergy(s))
println("Initial total angular momentum = ", TotalAngularMomentum(s))
println("Number of bodies = ", s.numberOfBodies)
for b in s.bodies
   println("body name = ", b.name)
end

t_final = 10.0 # final time of simulation
tspan = (0.0, t_final) # span of time to simulate
prob = ODEProblem(tendency!, s.phaseSpace, tspan, s) # 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("Final total energy = ", TotalEnergy(s))
println("Final total angular momentum  = ", TotalAngularMomentum(s))

body1 = 1
body2 = 2
# Plot of position vs. time
xt = plot(sample_times, sol[1, body1, :], xlabel = "t", ylabel = "x(t)", legend = false, title = "x vs. t")

# Plot of orbit
xy = plot([(sol[1, body1, :], sol[2, body1, :]), (sol[1, body2, :], sol[2, body2, :])], colors = ("yellow", "green"), xlabel = "x", ylabel = "y", legend = false, title = "Orbit")
#=
plot(xy, aspect_ratio=:equal)
=#

samples = 1000
interval = floor(Int,size(sol.t)[1] / samples)
animation = @animate for i=1:samples-1
   plot([(sol[1, body1, i*interval], sol[2, body1, i*interval]), (sol[1, body1, (i+1)*interval], sol[2, body1, (i+1)*interval])],
   aspect_ratio=:equal, colors = ("yellow", "green"), xlabel = "x", ylabel = "y", legend = false, title = "Orbit", xlims=(-1, 1), ylims = (-1, 1))
end

gif(animation, fps=15)