"""
Simple animation of how the Moon circles the Earth and the Earth-Moon-system orbits the Sun. 
"""



from itertools import count
import pandas as pd
import matplotlib
matplotlib.use("Qt4Agg")
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp, odeint
from matplotlib.animation import FuncAnimation
import numpy as np
import random

# Set up formatting for the movie files
Writer = matplotlib.animation.writers['ffmpeg']
writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)

Moon_mass_multiplier = 1
Earth_mass_multiplier = 1

G = 6.67408e-11 #m^3 /kg /s^2; gravitational constant
ms = 1.989e30   #kg; mass of Sun
me = 5.972e24   #kg; mass of Earth
mm = 7.3477e22  #kg true mass of moon
me *= Earth_mass_multiplier
mm *= Moon_mass_multiplier
Sun = 'Sun'
central_body = 'Earth'
orbiting_body = 'Moon'
if Moon_mass_multiplier != 1:
    orbiting_body = 'Moon x {:d}'.format(Moon_mass_multiplier)
d_Moon = 3.84402e8  #mean semi-major axis of the geocentric lunar orbit, measured from center to center.
d_Earth = 1.495978707e11  #astronomical unit
#wrong values:
#d_Moon = 3.84402e9  #mean semi-major axis of the geocentric lunar orbit, measured from center to center.
#d_Earth = 3.495978707e10  #astronomical unit

xs = np.array([0., 0.])  #initial position of Sun
xe = np.array([d_Earth, 0.])   #initial position of Earth
xm = np.array([d_Earth + d_Moon, 0.])   #initial position of Moon
v0s = 0.
v0e =  np.sqrt(G*ms/np.linalg.norm(xe-xs)) 
#v0m =  np.sqrt(G*np.linalg.norm(xm-xs) * (ms/np.linalg.norm(xm-xs)**2 + me/np.linalg.norm(xm-xe)**2))
v0m = v0e + np.sqrt(G*me/d_Moon)
vs = np.array([ 0., v0s])   #initial velocity of Earth
ve = np.array([ 0., v0e])   #initial velocity of Earth
vm = np.array([ 0., v0m])   #initial velocity of the Moon; m/s

print(xs, vs, xe, ve, xm, vm)
z0 = np.array([xs[0], vs[0], xe[0], ve[0], xm[0], vm[0], xs[1], vs[1], xe[1], ve[1], xm[1], vm[1]]) #initial conditions
print('initial conditions: {}'.format(z0))


def grav(z, t):
    dz = np.zeros_like(z)
    x1, vx1, x2, vx2, x3, vx3,  y1, vy1, y2, vy2, y3, vy3 = z #(initial conditions)
    dx1 = vx1
    dx2 = vx2
    dx3 = vx3
    dy1 = vy1
    dy2 = vy2
    dy3 = vy3
    dvx1 = -G*(me/np.linalg.norm(xs-xe)**3*(xs[0] - xe[0]) + mm/np.linalg.norm(xs-xm)**3*(xs[0] - xm[0]))
    dvx2 = -G*(ms/np.linalg.norm(xe-xs)**3*(xe[0] - xs[0]) + mm/np.linalg.norm(xe-xm)**3*(xe[0] - xm[0]))
    dvx3 = -G*(ms/np.linalg.norm(xm-xs)**3*(xm[0] - xs[0]) + me/np.linalg.norm(xm-xe)**3*(xm[0] - xe[0]))
    dvy1 = -G*(me/np.linalg.norm(xs-xe)**3*(xs[1] - xe[1]) + mm/np.linalg.norm(xs-xm)**3*(xs[1] - xm[1]))
    dvy2 = -G*(ms/np.linalg.norm(xe-xs)**3*(xe[1] - xs[1]) + mm/np.linalg.norm(xe-xm)**3*(xe[1] - xm[1]))
    dvy3 = -G*(ms/np.linalg.norm(xm-xs)**3*(xm[1] - xs[1]) + me/np.linalg.norm(xm-xe)**3*(xm[1] - xe[1]))
    dz[0] = dx1; dz[1] = dvx1; dz[2] = dx2; dz[3] = dvx2; dz[4] = dx3; dz[5] = dvx3; dz[6] = dy1
    dz[7] = dvy1; dz[8] = dy2; dz[9] = dvy2; dz[10] = dy3; dz[11] = dvy3
    return dz
    
#this is the part that does the orbit integration
t_end = 27.*86400. #seconds
n_steps = 150
t = np.linspace(0., t_end, n_steps)
sol = odeint(grav, z0, t)
xs = sol[:,0]; xe = sol[:,2]; xm = sol[:,4]; ys = sol[:,6]; ye = sol[:,8]; ym = sol[:,10]
print(sol[:,4], sol[:,10])

#set up the initial plt
fig, ax = plt.subplots()
d = 155.e9
ax.set_xlim(-d, d)
ax.set_ylim(-d, d)
line_s, = ax.plot(xs[0],xs[1])
line_e, = ax.plot(xe[0],xe[1])
line_m, = ax.plot(xm[0],xm[1])
ax.set_xlabel('x [m]',fontsize=14)
ax.set_ylabel('y [m]',fontsize=14)
ax.plot(xs[0], ys[0], 'k+', label = Sun)
ax.plot(xe[0], ye[0], 'b+', label = central_body)
ax.plot(xm[0], ym[0], 'r+', label = orbiting_body)
#ax.plot(R[0], R[1], 'b+')

x_vals_s = []
y_vals_s = []
x_vals_e = []
y_vals_e = []
x_vals_m = []
y_vals_m = []
index = count()

line_s, = ax.plot([], [], 'k.', markersize=0.2)
line_e, = ax.plot([], [], 'b.', markersize=0.2)
line_m, = ax.plot([], [], 'r.', markersize=2.)
pos_s, = ax.plot([], [], 'k.')
pos_e, = ax.plot([], [], 'b.')
pos_m, = ax.plot([], [], 'r.')


def animate(i):
    
    x_vals_s.append(xs[i])
    y_vals_s.append(ys[i])
    x_vals_e.append(xe[i])
    y_vals_e.append(ye[i])
    x_vals_m.append(xm[i])
    y_vals_m.append(ym[i])

    pos_s.set_xdata(xs[i])
    pos_s.set_ydata(ys[i])
    pos_e.set_xdata(xe[i])
    pos_e.set_ydata(ye[i])
    pos_m.set_xdata(xm[i])
    pos_m.set_ydata(ym[i])

    line_s.set_xdata(x_vals_s)
    line_s.set_ydata(y_vals_s)
    line_e.set_xdata(x_vals_e)
    line_e.set_ydata(y_vals_e)
    line_m.set_xdata(x_vals_m)
    line_m.set_ydata(y_vals_m)

    return line_s, line_e, line_m, pos_s, pos_e, pos_m
    
ani = FuncAnimation(fig, func=animate, frames=np.arange(0,n_steps-1,1), interval=1)

plt.tight_layout()
plt.legend(loc='best')
plt.gca().set_aspect('equal', adjustable='box')

try:
    writer = matplotlib.animation.writers['avconv']
except KeyError:
    writer = matplotlib.animation.writers['ffmpeg']
writer = writer(fps=3)
filename = 'Sun_Earth_Moon-animation.mp4'
if Moon_mass_multiplier != 1:
    filename = 'Sun_Earth_Moon-{:d}-animation.mp4'.format(Moon_mass_multiplier)

ani.save(filename, writer=writer, dpi=300)
plt.close(fig)