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

#plt.style.use('fivethirtyeight')

G = 6.67408e-11 #m^3 /kg /s^2; gravitational constant
m1 = 5.972e24   #kg; mass of Earth
m2 = 7.3477e22  #kg; mass of Moon
GMm = G*m1*m2
mu = m1*m2/(m1+m2) #reduced mass mu
central_body = 'Earth'
orbiting_body = 'Moon'
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
    

t_end = 29.*86400. #days
t = np.linspace(0. ,t_end, 1001)
sol = odeint(grav, z0, t)
rx = sol[:,0]; vx = sol[:,1]; ry = sol[:,2]; vy = sol[:,3]
xx = R[0] + m2/(m1+m2)*rx; xy = R[1] + m2/(m1+m2)*ry; yx = R[0] - m1/(m1+m2)*rx; yy = R[1] - m1/(m1+m2)*ry

fig, ax = subplots()
d = 450.e6
ax.set_xlim(-d, d)
ax.set_xlim(-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)

plt.plot(xx[0], xy[0], 'r+')
plt.plot(yx[0], yy[0], 'ko')
plt.plot(xx, xy, 'r-', label = central_body)
plt.plot(yx, yy, 'k', label = orbiting_body)
plt.legend(loc='best')
plt.gca().set_aspect('equal', adjustable='box')
plt.savefig('try-motion.pdf',bbox_inches='tight')
plt.show()


#x_vals = []
#y_vals = []

#index = count()

#fig, ax = plt.subplots()
#ax.set_xlim(0,100)
#ax.set_ylim(0,100)
#line, = ax.plot(0, 0)
#ax.set_xlabel('x label',fontsize=14)
#ax.set_ylabel('y label',fontsize=14)

#def animate(i):
#    x_vals.append(next(index))
#    y_vals.append(random.randint(0,5))

#    line.set_xdata(x_vals)
#    line.set_ydata(y_vals)
#    return line,
    
    #ax.cla()
    #ax.plot(x_vals, y_vals)


#ani = FuncAnimation(plt.gcf(), animate, interval=100)
#ani = FuncAnimation(fig, func=animate, frames=np.arange(0,10,1), interval=100)


#plt.tight_layout()
#plt.show()