#!/usr/bin/env python

from scipy.integrate import odeint
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from numpy.linalg import norm


############################################################
# Trabajamos en un sistema de unidades en el que G=1.0

G=1.0
m1=1.0
m2=0.001
mu=G*(m1+m2)


############################################################
# Vectores de condiciones iniciales y de tiempo

y0=[1.0, 0.0, 0.0, 1.2]
t=np.linspace(0.0,1000.0,1000.0)


############################################################
# Definicion de la funcion que devuelve la derivada temporal
# del vector de estado 'y'

def func(y,t):
    r=y[:2]
    v=y[2:]
    return [v[0], v[1], -mu*r[0]/norm(r)**3, -mu*r[1]/norm(r)**3]


############################################################
# Integracion con 'odeint'

solucion=odeint(func,y0,t)

r=solucion[:, :2]     # vector de posicion para cada t
v=solucion[:, 2:]     # vector de velocidad para cada t

############################################################
# Calculo de la energia especifica

Es=[]
for i,j in zip(r,v):
    Es+=[0.5*norm(j)**2-mu/norm(i)]

############################################################
# Graficacion

plt.figure()
plt.plot(t,Es)
plt.xlabel("t")
plt.ylabel("Energy")
plt.title("Energy vs time")
plt.savefig('EnergyvsTime.png')

plt.figure()
plt.subplot(121)
plt.plot(r[:,0],r[:,1])
plt.xlabel('x')
plt.ylabel('y')
plt.gca().set_aspect('equal', adjustable='box') # Equal size ratio
plt.subplot(122)
plt.plot(r[:,0],v[:,0])
plt.xlabel('x')
plt.ylabel('v_x')
plt.savefig('SpaceAndPhase.png')