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

# Simulate a solar system

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

println("Number of threads = ", Threads.nthreads())

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
   v::Vector{Float64} # velocity vector
end 

function star(name="Sun", m = 1.0, r = zeros(3), v = zeros(3))
   return body(name, m, r, v)
end

function planet(name="Earth", m=3.0e-6, a=1.0, ϵ=0.017, i=0.0)
   perihelion = (1.0 - ϵ) * a
	aphelion = (1.0 + ϵ) * a
	speed = sqrt(G * (1.0 + ϵ)^2 / (a * (1.0 - ϵ^2)))
   phi = 0.0
	r = [perihelion * cos(i*pi/180.0) * cos(phi), perihelion * cos(i*pi/180.0) * sin(phi), perihelion * sin(i*pi/180.0)]
	v = [-speed * sin(phi), speed * cos(phi), 0.0]
   return body(name, m, r, v)
end

function momentum(b::body)
   return b.m * b.v
end

function angularMomentum(b::body)
   return b.m * cross(b.r, b.v)
end

function kineticEnergy(b::body)
   v = b.v
   m = b.m
   return 0.5 * m * dot(v, v)
end

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

function rv(b::body)
   return [b.r; b.v]
end

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

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

function SolarSystem()

   bodies = Vector{body}()
   push!(bodies, star()) # default = Sun
   # push!(bodies, planet("Venus", 2.44e-6, 0.72, 0.0068, 3.39))
   # push!(bodies, planet("Earth", 3.0e-6, 1.0, 0.017, 0.0))
   push!(bodies, planet("Jupiter", 0.00095, 5.2, 0.049, 1.3))
   # push!(bodies, planet("Saturn", 0.000285, 9.58, 0.055, 2.48))
   push!(bodies, body("Astroid 1", 1e-13, [3.0, 0.0, 0.0], [0.0, 3.628, 0.0]))
   push!(bodies, body("Astroid 2 ", 1e-13, [3.276, 0.0, 0.0], [0.0, 3.471, 0.0]))
   # push!(bodies, body("Astroid 3 (no gap)", 1e-13,[3.7, 0.0, 0.0], [0.0, 3.267, 0.0]))
   numberOfBodies = size(bodies)[1]

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

end

function TotalMass(s::SolarSystem)
   M = 0.0
   for b in s.bodies
      M += b.m
   end
   return M
end

function TotalLinearMomentum(s::SolarSystem)
   P = zeros(3)
   for b in s.bodies
      P += momentum(b)
   end
   return P
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 CenterOfMassFrame!(s::SolarSystem)

   M = TotalMass(s)
   P = TotalLinearMomentum(s)
   V = P / M
   
   s.phaseSpace = zeros(6, 0)
   for body in s.bodies
      body.v -= V #boost to COM frame
      s.phaseSpace = [s.phaseSpace rv(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.v = 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.v  #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] += b2.m * f
            dps[4:6, j] -= b1.m * f # Newton's 3rd law
      end
   end

   return nothing
end

function parallel_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.v = ps[4:6, i]
      i += 1
   end

   # find velocities of bodies and forces on them. O(N^2) computational cost
   N = s.numberOfBodies
   Threads.@threads for i in 1:N
      b1 = s.bodies[i]
      dps[1:3, i] = b1.v  #dr/dt
      dps[4:6, i] = zeros(3)
      for j in 1:N
            if (j != i)
               b2 = s.bodies[j]
               f = force(b1, b2)
               dps[4:6, i] += b2.m * f
            end
      end
   end

   return nothing
end

s = SolarSystem()
CenterOfMassFrame!(s)
println(typeof(s))
println("Initial total energy = ", TotalEnergy(s))
println("Initial total linear momentum = ", TotalLinearMomentum(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 = 200.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, maxiters=1e8, reltol=1e-10, abstol=1e-10) # solve using Tsit5 algorithm to specified accuracy

println("\n\t Results")
println("Final time  = ", sol.t[end])
println("Final total energy = ", TotalEnergy(s))
println("Final total linear momentum = ", TotalLinearMomentum(s))
println("Final total angular momentum  = ", TotalAngularMomentum(s))

function eccentricity(body::Int64, solution)
   n = size(solution[1, 1, :])[1]
   samples = 1000
   interval = floor(Int, n/samples)
   ϵ = zeros(samples-1)
   time = zeros(samples-1)
   for i in 1:samples-1
      r2 = sol[1, body, i*interval:(i+1)*interval].^2 + sol[2, body, i*interval:(i+1)*interval].^2 + sol[3, body, i*interval:(i+1)*interval].^2
      aphelion = sqrt(maximum(r2))
      perihelion = sqrt(minimum(r2))
      ϵ[i] = (aphelion - perihelion) / (aphelion + perihelion)
      time[i] = solution.t[i*interval]
   end
   return (time, ϵ)
end

function obliquity(body::Int64, solution)
   n = size(solution[1, 1, :])[1]
   samples = 1000
   interval = floor(Int, n/samples)
   ob = zeros(samples)
   time = zeros(samples)
   for i in 1:samples
      r = solution[1:3, body, i*interval]
      v = solution[4:6, body, i*interval]
      ell = cross(r, v)
      norm = sqrt(dot(ell, ell))
      ob[i] = (180.0 / pi) * acos(ell[3]/norm)
      time[i] = solution.t[i*interval]
   end
   return (time, ob)
end


body1 = 2 # Jupiter
body2 = 3 # Astroid 1
body3 = 4 # Astroid 2
# body4 = 5 # Astroid 3

# Plot of orbit
xy = scatter([(sol[1, body1, :], sol[2, body1, :]), (sol[1, body2, :], sol[2, body2, :]),
 (sol[1, body3, :], sol[2, body3, :]), 
 # (sol[1, body4, :], sol[2, body4, :])
 ], 
 xlabel = "x (AU)", ylabel = "y (AU)", title = "Orbit",
 label = ["Jupiter" "Asteroid not in gap (3.0AU)" "Asteroid 2/1 gap (3.27AU)" "Asteroid not in gap (3.7AU)"],
 aspect_ratio=:equal, markersize=.5, markerstrokewidth = 0, legend = :bottomright,
 colors=[:red :blue :green :black]
 )

# Plot of obliquity
tilt = obliquity(body2, sol)
obliquityPlot = plot(tilt, xlabel = "t", ylabel = "tilt", legend = false, title = "obliquity")

# Plot of eccentricity
eccentric = eccentricity(body2, sol)
eccentricityPlot = plot(eccentric, xlabel = "t", ylabel = "ϵ", legend = false, title = "eccentricity")

#plot(obliquityPlot, eccentricityPlot)
plot(xy, aspect_ratio=:equal)

savefig("4-17.png")