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

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))

G    = 6.67408e-08            # Gravitational constant in CGS units [cm^3/g/s^2]
Msun = 1.98892e33             # Solar mass              [g]
Mju  = 1.899e30               # Mass of Jupiter         [g]
year = 31557600.e0            # Year                    [s]
au   = 1.49598e13             # Astronomical Unit       [cm]
t0   = 0.                     # Starting time
tend = 4.5*year               # End time
nt   = 400                    # Nr of time steps
t    = np.linspace(t0,tend,nt)
m    = np.array([Msun,Mju,Mju])   # Masses
a0   = np.array([0.,au,1.15*au,]) # Semi-major axes
phi0 = np.array([0.,0.,0.])       # Orbital locations in degrees
nb   = len(m)
x0   = np.zeros((nb,3))
v0   = np.zeros((nb,3))
pos  = np.zeros(3)
mom  = np.zeros(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.)
    pos[:] += m[i]*x0[i,:]
    mom[:] += m[i]*v0[i,:]
x0[0,:] = -pos[:]/m[0]        # Total center of mass must be at (0,0,0)
v0[0,:] = -mom[:]/m[0]        # Total momentum must be (0,0,0)
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]
plt.figure()
for ibody in range(nb):
    plt.plot(x[ibody,0,:]/au,x[ibody,1,:]/au)
plt.xlabel('x [au]')
plt.ylabel('y [au]')
plt.axis('equal')
plt.savefig('fig_nbody_1_1.pdf')
plt.show()