import numpy as np
import matplotlib.pyplot as plt

# Constants
m1 = 5.972E24  # kg mass Earth
m2 = 4.19725E5 # kg mass ISS
mu = m1*m2/(m1+m2)
G = 6.674E-11  # m^3 kg^-1 s^-2  gravitational constant
l = 2.20E16     # kg m^2 s^-1 angular momentum
#l = 2.12E16     # kg m^2 s^-1 angular momentum
rE = 6.371E6   # m radius of Earth
p = 5.0E9        # kg m s^-1 typical radial momentum

# Make plotting grid
rlist  = np.linspace(0.3*rE, 1.8*rE, 500)
prlist = np.linspace(-p, p, 500)
r, pr = np.meshgrid(rlist, prlist)

# Derivatives
rdot = pr/mu
prdot = l**2/mu/r**3-G*m1*m2/r**2

# Plot phase space using streamplot
# streamplot uses (x,y) coordinates and their deriatives as input
plt.streamplot(r, pr, rdot, prdot, density=2)

# Plot minimum of Hamiltonian
plt.plot([l**2/mu/(G*m1*m2)], [0], 'rx')
print 'Minimum energy for r=', l**2/mu/(G*m1*m2)

# Plot zero energy
#pr=np.sqrt(2*mu*G*m1*m2/r-l**2/r**2)
rE0 = np.linspace(0.3*rE, 1.8*rE, 10000)
prE0 = np.sqrt(2*mu*G*m1*m2/rE0-l**2/rE0**2) # Finds values of p_r giving H=0
plt.plot(rE0, prE0, 'r')
plt.plot(rE0, -prE0, 'r')

# Plot radius of Earth
rEarth = [rE,rE] # m
prEarth = [-p,p]
plt.plot(rEarth, prEarth, color='green', linestyle='dashed', linewidth=2)

# Decorations
plt.xlabel('$r$ [m]')
plt.ylabel('$p_r$ [kg m s$^{-1}$]')
plt.xlim([0.3*rE, 1.8*rE])
plt.ylim([-p, p])
plt.savefig('phase_space.png')