import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
import matplotlib.animation as animation

def f(y,t,m):
    G    = 6.67408e-08    # Gravitational constant in CGS units [cm^3/g/s^2]
    n    = len(m)
    plan = y.reshape((n,6))
    frc  = np.zeros((n,3))
    dy   = np.zeros((n,6))
    for i in range(n):
        x    = plan[i,0:3]
        for j in range(i+1,n):
            x1    = plan[j,0:3]
            dx    = x1-x
            r     = np.sqrt(dx[0]**2+dx[1]**2+dx[2]**2)
            df    = G*m[i]*m[j]*dx/r**3
            frc[i,0:3] += df[0:3]
            frc[j,0:3] -= df[0:3]
    for i in range(n):
        x    = plan[i,0:3].copy()
        v    = plan[i,3:6].copy()
        dxdt = v
        dvdt = frc[i,0:3]/m[i]
        dy[i,0:6] = np.hstack((dxdt,dvdt))
    return dy.reshape((6*n))

def center(m,x0,v0):
    nb   = len(m)
    pos  = np.zeros(3)
    for i in range(nb):
        pos[:] += m[i]*x0[i,:]
    xnew = np.zeros_like(x0)
    vnew = np.zeros_like(v0)
    mtot = m.sum()
    for i in range(nb):
        xnew[i,:] = x0[i,:] - pos[:] / mtot
    mom  = np.zeros(3)
    for i in range(nb):
        mom[:] += m[i]*v0[i,:]
    for i in range(nb):
        vnew[i,:] = v0[i,:] - mom[:] / mtot
    return xnew,vnew
    
def nbodyint(m,x0,v0,t):
    xv0 = np.zeros((nb,6))        
    for i in range(nb):
        xv0[i,0:3] = x0[i,0:3]
        xv0[i,3:6] = v0[i,0:3]
    y0   = xv0.reshape(6*nb)         # The full nb*6 element vector
    sol  = odeint(f,y0,t,args=(m,))  # Solve the N-body problem
    xv   = sol.reshape((nt,nb,6))    # Now extract again the x and v 
    x    = np.zeros((nb,3,nt))
    v    = np.zeros((nb,3,nt))
    for i in range(nb):
        for idir in range(3):
            x[i,idir,:] = xv[:,i,idir]
            v[i,idir,:] = xv[:,i,3+idir]
    return x,v

G    = 6.67408e-08            # Gravitational constant in CGS units [cm^3/g/s^2]
Msun = 1.98892e33             # Solar mass              [g]
Mea  = 5.9736e27              # Mass of Earth           [g]
year = 31557600.e0            # Year                    [s]
au   = 1.49598e13             # Astronomical Unit       [cm]
t0   = 0.                     # Starting time
tend = 0.04*year               # End time
nt   = 100                   # Nr of time steps
t    = np.linspace(t0,tend,nt)
m    = np.array([0.0802*Msun,0.85*Mea,1.38*Mea,0.41*Mea,0.62*Mea,0.68*Mea,1.34*Mea,0.01*Mea])    # Masses (last one taken small)
a0   = np.array([0.,11.11e-3*au,15.21e-3*au,21.44e-3*au,28.17e-3*au,37.1e-3*au,45.1e-3*au,63e-3*au])  # Semi-major axes
phi0 = np.array([0.,260.,0.,125.,40.,285,15.,150.])      # Orbital angular locations in degrees (estim from Fig. 1 Gillon etal)
nb   = len(m)
x0   = np.zeros((nb,3))
v0   = np.zeros((nb,3))

for i in range(1,nb):         # Loop over all planets (i.e. excluding star)
    x0[i,0] = a0[i]*np.cos(phi0[i]*np.pi/180.)
    x0[i,1] = a0[i]*np.sin(phi0[i]*np.pi/180.)
    vphi    = np.sqrt(G*m[0]/a0[i])
    v0[i,0] = -vphi*np.sin(phi0[i]*np.pi/180.)
    v0[i,1] =  vphi*np.cos(phi0[i]*np.pi/180.)

x0,v0 = center(m,x0,v0)
x,v   = nbodyint(m,x0,v0,t)

count = 0

def update(frameNum, a0):
    global count, nb, nt, x, v, t
    a0.set_data(x[:,0,count]/au,x[:,1,count]/au)
    count = (count + 1) % nt

rmau= 0.07
fig = plt.figure()
ax  = plt.axes(xlim=(-rmau, rmau), ylim=(-rmau, rmau))
ax.set(aspect='equal',xlabel='x [au]',ylabel='y [au]')
a0, = ax.plot(x[:,0,0]/au,x[:,1,0]/au,'o')

anim = animation.FuncAnimation(fig,update,fargs=(a0,),interval=120)
plt.show()