"""
Simple animation of how the Moon circles the Earth. One can barely see that the Earth in fact circles around the common senter of mass.
Increasing the "moon's" mass by a factor of 10 shows this more clearly.
In the situation where the moon's mass is 100 times that of the real moon, one clearly sees a qualittively different behavior. 
"""



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

G = 6.67408e-11 #m^3 /kg /s^2; gravitational constant
m1 = 5.972e24   #kg; mass of Earth
m2 = 7.3477e22  #kg true mass of moon
m2 *= Moon_mass_multiplier
GMm = G*m1*m2
mu = m1*m2/(m1+m2) #reduced mass mu
central_body = 'Earth'
orbiting_body = 'Moon'
if Moon_mass_multiplier != 1:
    orbiting_body = 'Moon x {:d}'.format(Moon_mass_multiplier)
d_moon = 384402.e3  #mean semi-major axis of the geocentric lunar orbit, measured from center to center.

x1 = np.array([0., 0.])   #initial position of Earth
x2 = np.array([0., d_moon])   #initial position of Moon
R = np.array((m1*x1 + m2*x2)/(m1 + m2)) #distance of center of mass from Earth's center
print('distance of center of mass from center of Earth = {}'.format(R))
v0m1 = 0.
v0m2 =  np.sqrt(G*m1/np.linalg.norm(x2-x1)) 
v1 = np.array([ v0m1, 0.])   #initial velocity of Earth
v2 = np.array([-v0m2, 0.])   #initial velocity of the Moon; m/s
z0 = np.array([x1[0] - x2[0], v1[0] - v2[0], x1[1] - x2[1], v1[1] - v2[1] ]) #initial conditions
print('initial conditions: {}'.format(z0))


def grav(z, t):
    dz = np.zeros_like(z)
    rx, vx, ry, vy = z #(initial conditions)
    r3 = np.linalg.norm([rx,ry])**3
    drx = vx
    dry = vy
    dvx = -GMm*rx/r3/mu
    dvy = -GMm*ry/r3/mu
    dz[0] = drx; dz[1] = dvx; dz[2] = dry; dz[3] = dvy
    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)
rx = sol[:,0]; vx = sol[:,1]; ry = sol[:,2]; vy = sol[:,3]
x1 = R[0] + m2/(m1+m2)*rx; y1 = R[1] + m2/(m1+m2)*ry; x2 = R[0] - m1/(m1+m2)*rx; y2 = R[1] - m1/(m1+m2)*ry

#set up the initial plt
fig, ax = plt.subplots()
d = 450.e6
ax.set_xlim(-d, d)
ax.set_ylim(-d, d)
line_e, = ax.plot(x1[0],x1[1])
line_m, = ax.plot(x2[0],x2[1])
ax.set_xlabel('x [m]',fontsize=14)
ax.set_ylabel('y [m]',fontsize=14)
ax.plot(x1[0], y1[0], 'r+', label = central_body)
ax.plot(x2[0], y2[0], 'k+', label = orbiting_body)
ax.plot(R[0], R[1], 'b+')

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

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


def animate(i):
    
    x_vals_e.append(x1[i])
    y_vals_e.append(y1[i])
    x_vals_m.append(x2[i])
    y_vals_m.append(y2[i])

    pos_e.set_xdata(x1[i])
    pos_e.set_ydata(y1[i])
    pos_m.set_xdata(x2[i])
    pos_m.set_ydata(y2[i])

    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_e, line_m, 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 = 'Earth_Moon-animation.mp4'
if Moon_mass_multiplier != 1:
    filename = 'Earth_Moon-{:d}-animation.mp4'.format(Moon_mass_multiplier)

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