import numpy as np
import matplotlib.pyplot as plt

object = str(input("Would you like to calculate the orbit for Earth or Halley's Comet? (E/H): "))

au_m = 149597870700

G = 6.673e-11       # Gravitational constant [N (m/kg)^2]
M = 1.9891e30       # Solar mass [kg]

if (object=='E' or object=='e'):
    a = 1.000*au_m
    epsilon = 0.017
    m = 5.9763e24
    aphelion = a*(1+epsilon)/au_m
    perihelion = a*(1-epsilon)/au_m
    R = a*(1-epsilon)
    T = (2*np.pi)*np.sqrt((a**3)/G/M)
    E = G*M*m/2/a
    #V0 = np.sqrt(G*(M+m)*((2/R)-(1/a)))*T/R
elif (object=='H' or object=='h'):
    a = 17.95*au_m
    epsilon = 0.967
    m = 1e14
    aphelion = a*(1+epsilon)/au_m
    perihelion = a*(1-epsilon)/au_m
    R = a*(1-epsilon)
    T = (2*np.pi)*np.sqrt((a**3)/G/M)
    E = G*M*m/2/a
    #V0 = np.sqrt(G*(M+m)*((2/R)-(1/a)))*T/R
    Vcross = np.sqrt(G*(M+m)*((2/(1.000*au_m))-(1/a)))/1e3


C = (G*M*T**2)/(R**3)

# Initial conditions
X0 = 1
U0 = 0
Y0 = 0
V0 = np.sqrt(G*(M+m)*((2/R)-(1/a)))*T/R


# Get number of steps per orbit
tau = 5
N = int(input("How many steps would you like to calculate per orbit? "))*tau
dt = tau/(N)
t = np.linspace(0,tau,N)*T/24/3600




# Define Runge-Kutta functions
def F1(X_, Y_, U_, V_): # dX/dtau
    return U_

def F2(X_, Y_, U_, V_): # dY/dtau
    return V_

def F3(X_, Y_, U_, V_): # dU/dtau
    return -C * (  X_/( (np.sqrt(X_**2 + Y_**2) )**3)  )

def F4(X_, Y_, U_, V_): # dV/dtau
    return -C * (  Y_/( (np.sqrt(X_**2 + Y_**2) )**3)  )


# Initialize solution arrays
X_4RK = np.zeros(N)
Y_4RK = np.zeros(N)
U_4RK = np.zeros(N)
V_4RK = np.zeros(N)

# Set initial position and velocity in the x array and y-velocity array
X_4RK[0] = X0
V_4RK[0] = V0

# Perform 4th order Runge-Kutta
for n in range(N-1):

    k_x1 = dt * F1( X_4RK[n], Y_4RK[n], U_4RK[n], V_4RK[n] )
    k_y1 = dt * F2( X_4RK[n], Y_4RK[n], U_4RK[n], V_4RK[n] )
    k_u1 = dt * F3( X_4RK[n], Y_4RK[n], U_4RK[n], V_4RK[n] )
    k_v1 = dt * F4( X_4RK[n], Y_4RK[n], U_4RK[n], V_4RK[n] )

    k_x2 = dt * F1( X_4RK[n] + k_x1/2, Y_4RK[n] + k_y1/2, U_4RK[n] + k_u1/2, V_4RK[n] + k_v1/2 )
    k_y2 = dt * F2( X_4RK[n] + k_x1/2, Y_4RK[n] + k_y1/2, U_4RK[n] + k_u1/2, V_4RK[n] + k_v1/2 )
    k_u2 = dt * F3( X_4RK[n] + k_x1/2, Y_4RK[n] + k_y1/2, U_4RK[n] + k_u1/2, V_4RK[n] + k_v1/2 )
    k_v2 = dt * F4( X_4RK[n] + k_x1/2, Y_4RK[n] + k_y1/2, U_4RK[n] + k_u1/2, V_4RK[n] + k_v1/2 )

    k_x3 = dt * F1( X_4RK[n] + k_x2/2, Y_4RK[n] + k_y2/2, U_4RK[n] + k_u2/2, V_4RK[n] + k_v2/2 )
    k_y3 = dt * F2( X_4RK[n] + k_x2/2, Y_4RK[n] + k_y2/2, U_4RK[n] + k_u2/2, V_4RK[n] + k_v2/2 )
    k_u3 = dt * F3( X_4RK[n] + k_x2/2, Y_4RK[n] + k_y2/2, U_4RK[n] + k_u2/2, V_4RK[n] + k_v2/2 )
    k_v3 = dt * F4( X_4RK[n] + k_x2/2, Y_4RK[n] + k_y2/2, U_4RK[n] + k_u2/2, V_4RK[n] + k_v2/2 )

    k_x4 = dt * F1( X_4RK[n] + k_x3, Y_4RK[n] + k_y3, U_4RK[n] + k_u3, V_4RK[n] + k_v3 )
    k_y4 = dt * F2( X_4RK[n] + k_x3, Y_4RK[n] + k_y3, U_4RK[n] + k_u3, V_4RK[n] + k_v3 )
    k_u4 = dt * F3( X_4RK[n] + k_x3, Y_4RK[n] + k_y3, U_4RK[n] + k_u3, V_4RK[n] + k_v3 )
    k_v4 = dt * F4( X_4RK[n] + k_x3, Y_4RK[n] + k_y3, U_4RK[n] + k_u3, V_4RK[n] + k_v3 )

    X_4RK[n+1] = (X_4RK[n] + k_x1/6 + k_x2/3 + k_x3/3 + k_x4/6)
    Y_4RK[n+1] = (Y_4RK[n] + k_y1/6 + k_y2/3 + k_y3/3 + k_y4/6)
    U_4RK[n+1] = U_4RK[n] + k_u1/6 + k_u2/3 + k_u3/3 + k_u4/6
    V_4RK[n+1] = V_4RK[n] + k_v1/6 + k_v2/3 + k_v3/3 + k_v4/6


# Is the orbit closed?
initial = np.sqrt((X_4RK[0]*au_m)**2+(Y_4RK[0]*au_m)**2)
final = np.sqrt((X_4RK[N-1]*au_m)**2+(Y_4RK[N-1]*au_m)**2)
if (abs(initial-final)<initial*(0.1/100)):
    co = str('The orbit is closed.')
else:
    co = str('The orbit is not closed.')



# Format information for plotting
deltat = str('dt ='+'{:.4f}'.format(dt*T/24/3600)+'days')
stepsperorbit = str('N(steps/orbit) = '+'{:.0f}'.format(N/tau))
Norbit = str('N(orbit) = '+'{:.0f}'.format(tau))
energy = str('E = '+'{:.3e}'.format(E)+' Joules')

# Put position solutions in units of meters
x = X_4RK*R
y = Y_4RK*R


plt.figure(figsize=(16,5))
#plt.plot(X_4RK[0:int(9*N/10)], Y_4RK[0:int(9*N/10)], 'g', [0], [0], 'ro', [X_4RK[len(X_4RK)-1]], [Y_4RK[len(Y_4RK)-1]], 'bo')
#plt.plot(X_4RK*au_m, Y_4RK*au_m, 'g', [0], [0], 'ro', [X_4RK[len(X_4RK)-1]*au_m], [Y_4RK[len(Y_4RK)-1]*au_m], 'bo')
if (object=='E' or object=='e'):
    plt.title("Earth's Solar Orbit", fontsize=20)
    plt.axvline(x=0, color='k')
    plt.axhline(y=0, color='k')
    plt.plot(x/au_m, y/au_m, 'g', [x[len(X_4RK)-1]/au_m], [y[len(Y_4RK)-1]/au_m], 'bo')
    plt.xlabel(r"$x$ (AU)", fontsize=20)
    plt.ylabel(r"$y$ (AU)", fontsize=20)
    plt.plot(0,0,'.', markersize=30, color='yellow')
    plt.text(0.1,0.7, f"Perihelion = {'{:.2f}'.format(perihelion)} AU")
    plt.text(0.1,0.6, f"Aphelion = {'{:.2f}'.format(aphelion)} AU")
    plt.text(0.1,0.5, f"Orbital Period = {'{:.2f}'.format(T/365/24/3600)} Years ")
    plt.text(0.1,0.4, stepsperorbit)
    plt.text(0.1,0.3, Norbit)
    plt.text(0.1,0.2, deltat)
    plt.text(0.1,0.1, energy)
    plt.text(0.1, -0.5, co, fontsize=20)
elif (object=='H' or object=='h'):
    plt.title("Halley's Solar Orbit", fontsize=20)
    plt.axvline(x=0, color='k')
    plt.axhline(y=0, color='k')
    plt.plot(x/au_m, y/au_m, 'g', [x[len(X_4RK)-1]/au_m], [y[len(Y_4RK)-1]/au_m], 'bo')
    plt.xlabel(r"$x$ (AU)", fontsize=20)
    plt.ylabel(r"$y$ (AU)", fontsize=20)
    plt.plot(0,0,'.', markersize=30, color='yellow')
    plt.text(-25,3, f"Perihelion = {'{:.2f}'.format(perihelion)} AU")
    plt.text(-25,2.6, f"Aphelion = {'{:.2f}'.format(aphelion)} AU")
    plt.text(-25,2.2, f"Orbital Period = {'{:.2f}'.format(T/365/24/3600)} Years ")
    plt.text(-25,1.8, stepsperorbit)
    plt.text(-25,1.4, Norbit)
    plt.text(-25,1.0, deltat)
    plt.text(-25,0.6, energy)
    plt.text(-25, -2, co, fontsize=20)

plt.savefig('Orbit.eps', dpi=500)
plt.show()

#d = np.sqrt(X_4RK**2+Y_4RK**2)*R

# Calculate total distance from the Sun
d = np.sqrt(x**2+y**2)/au_m
plt.figure(figsize=(16,5))
if (object=='E' or object=='e'):
    plt.title("Earth's Distance from Sun", fontsize=20)
elif (object=='H' or object=='h'):
    plt.title("Halley's Distance from Sun", fontsize=20)
plt.plot(t/365, d, 'b')
plt.ylabel("Distance from Sun (AU)", fontsize=20)
plt.xlabel(r"Time (Years)", fontsize=20)

plt.savefig('dist_time.eps', dpi=500)
plt.show()


# Calculate total velocity
vtotal = np.sqrt(U_4RK**2+V_4RK**2)*R/T/1e3
plt.figure(figsize=(16,5))
plt.plot(d,vtotal, 'k')
plt.xlabel('Distance from Sun (AU)', fontsize=20)
plt.ylabel('Orbital Velocity (km/s)', fontsize=20)
if (object=='H' or object=='h'):
    plt.title("Halley's Orbital Velocity vs Distance", fontsize=20)
    plt.plot(1.000, Vcross, 'r.', markersize=10)
    plt.text(1.000+0.2*1.000, Vcross,f"Velocity at Earth's orbit = {'{:.2f}'.format(Vcross)} km/s")
elif (object=='E' or object=='e'):
    plt.title("Earth's Orbital Velocity vs Distance", fontsize=20)

plt.savefig('vel_dist.eps', dpi=500)
plt.show()


# Calculate angle based on x,y position
theta = np.empty(int(N/tau))
for i in range(0,int(N/tau)):
    if (y[i]>0):
        theta[i] = np.arccos(x[i]/(np.sqrt((x[i])**2+(y[i])**2)))*180/np.pi
    else:
        theta[i] = 360-np.arccos(x[i]/(np.sqrt((x[i])**2+(y[i])**2)))*180/np.pi


plt.figure(figsize=(16,5))
plt.plot(theta, vtotal[0:int(N/tau)], 'k')
#plt.plot(theta, vtotal, 'g')
plt.xlabel(r'$\theta$ (degrees)', fontsize=20)
plt.ylabel(r'Orbital Velocity (km/s)', fontsize=20)
if (object=='E' or object=='e'):
    plt.title("Earth's Orbital Velocity vs Angle")
elif (object=='H' or object=='h'):
    plt.title("Halley's Orbital Velocity vs Angle")

plt.savefig('vel_angle.eps', dpi=500)
plt.show()