from numpy import sqrt, array, cos, sin, pi, arange, array
from pylab import plot, show, legend, xlabel, ylabel
%matplotlib inline

G = 6.674e-11
M = 5.972e24
m = 7.348e22 # mass of the moon

def alpha(r):
    radiusMagnitude = sqrt(x**2 + y**2)
    return -G * M * r / (radiusMagnitude**3)

# This program solves dr/dt = f(r,t)
def f(r,t):

    # unpack variables
    x = r[0]
    y = r[1]
    vx = r[2]
    vy = r[3]

    radiusMagnitude = sqrt(x**2 + y**2)

    fx = vx
    fy = vy
    fvx = -G * M * x / (radiusMagnitude**3)
    fvy = -G * M * y / (radiusMagnitude**3)

    return array([fx, fy, fvx, fvy], float)

# initial conditions
t0 = 0.0                          # initial time
tf = 6.998e6 # 81 days
h = 1 # time step of one hour
tpoints = arange(t0, tf, h) # array of time values

x = -3.626e8                          # initial x in meters
y = 0                          # intiial y
vx = 0                # initial vx
vy = 1.077e3       # initial vy m/s
r = array([x, y, vx, vy], float) # initial r vector
xpointsV = []
ypointsV = []
vxpointsV = []
vypointsV = []
PElist = [] # potential energy
KElist = [] # kinetic energy
TElist = [] # total energy

vxmid = vx + alpha(x) * 0.5 * h
vymid = vy + alpha(y) * 0.5 * h

for t in tpoints:

    #verlet method
    xpointsV.append(x)
    ypointsV.append(y)
    vxpointsV.append(vx)
    vypointsV.append(vy)

    potential = -G * M * m / (sqrt(x**2 + y**2))
    kinetic = .5 * m * (vx**2 + vy**2)
    total = potential + kinetic
    PElist.append(potential)
    KElist.append(kinetic)
    TElist.append(total)

    x += vxmid * h
    y += vymid * h
    vx = vxmid + alpha(x) * 0.5 * h
    vy = vymid + alpha(y) * 0.5 * h
    vxmid += alpha(x) * h
    vymid += alpha(y) * h

plot (tpoints, PElist, "r-", label="potential")
plot (tpoints, KElist, "g-", label="kinetic")
plot (tpoints, TElist, "b-", label="total")
xlabel("time")
ylabel("energy (J)")
show()