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

%matplotlib inline

start_time = time.time()

M = 1.989e30
G = 6.674e-11

# 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 = 40 * 3.154e7 # 3.154e7 is # of seconds in years
N = 75000
# h = (tf - t0) / N # timestep
h = 5e4

x = 4e12                          # initial x in meters
y = 0.0                          # intiial y
vx = 0                # initial vx
vy = 500       # initial vy m/s
r = array([x, y, vx, vy], float) # initial r vector
xpoints = [x]
ypoints = [y]
hlist = [h]
tlist = [t0]

numcross = 0             # number of x = 0 crossings from left to right
delta = 1000 / 3.154e7             # desired accuracy of 1 km / year

while numcross <= 2:
    # adaptive step size

    # do one large step
    k1 = 2*h*f(r, t)
    k2 = 2*h*f(r + 0.5*k1, t + 0.5 * h)
    k3 = 2*h*f(r + 0.5*k2, t + 0.5 * h)
    k4 = 2*h*f(r + k3, t + h)
    r1 = r + (k1 + 2*k2 + 2*k3 + k4)/6

    # do two small steps
    k1 = h*f(r, t)
    k2 = h*f(r + 0.5*k1, t + 0.5 * h)
    k3 = h*f(r + 0.5*k2, t + 0.5 * h)
    k4 = h*f(r + k3, t + h)
    r2 = r + (k1 + 2*k2 + 2*k3 + k4)/6

    k1 = h*f(r2, t)
    k2 = h*f(r2 + 0.5*k1, t + 0.5 * h)
    k3 = h*f(r2 + 0.5*k2, t + 0.5 * h)
    k4 = h*f(r2 + k3, t + h)
    r2 += (k1 + 2*k2 + 2*k3 + k4)/6

    # calculate the value of rho
    x1,y1 = r1[0],r1[2]
    x2,y2 = r2[0],r2[2]
    epsilon = sqrt((x1 - x2)**2 + (y1 - y2)**2)
    rho = 30*h*delta/epsilon

    # if rho >= 1.0, increase step size, save data
    if rho >= 1.0:
        t += 2*h                  # increase t
        h *= min(rho**0.25, 2.0)  # increase h

        yold = r[1]
        r = r2                    # update r
        ynew = r[1]

        # test to see if object crossed the x-axis from negative to positive
        if ynew > 0 and yold < 0:
            numcross += 1

        # save data
        xpoints.append(r[0])
        ypoints.append(r[1])
        hlist.append(h)
        tlist.append(t)

    # if rho <= 1.0, decrease step size, repeat
    else:
        h *= rho**0.25

plot (ypoints, xpoints, "b.", label="y vs x")
xlabel("y pos")
ylabel("x pos")
legend()
show()

print("--- %s seconds taken to run ---" % (time.time() - start_time))