from math import *
import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt


def part1(velocityFactor, orbits):
    #Part one Diffeq
    def f_part1(t, state, Me, Mm, G):
        xm, ym, vx, vy = state #all input values

        dxmdt = vx #functions for each diffeq
        dymdt = vy
        dvxdt = -(Me*G*xm)/((xm**2+ym**2)**(3/2))
        dvydt = -(Me*G*ym)/((xm**2+ym**2)**(3/2))
        

        return (dxmdt, dymdt, dvxdt, dvydt) #return differentiated values


    # Initial Conditions
    Me = 5.97*(10**24) #Constants
    Mm = 7.35*(10**22)
    G = 6.67*(10**-11)

    t_min = 0
    t_max = 2360620*int(orbits) #Time for one orbit * number of orbits
    numpoints = 2000*int(orbits) #To keep good clarity when plotting
    t = np.linspace(t_min, t_max, numpoints)

    r = 384400000 #Distance or orbit
    v = sqrt((G*Me)/r)*float(velocityFactor)#Balancing centrepetal force and gravitational force,
                                            #then multiplying by a scalar to obtain an elliptical orbit

    xm0 = r #Initial position and velocity of moon
    ym0 = 0 #Set so moon is travelling only in the y direction initially
    vx0 = 0
    vy0 = v

    rtol = 1e-5 #Tolerences set to balance computation speed with accuracy
    atol = 1e-8


    #Solver
    results = solve_ivp(f_part1, (t_min,t_max), (xm0, ym0, vx0, vy0), args=(Me, Mm, G), t_eval=t, atol=atol, rtol=rtol)


    #Graph plotting
    ax=plt.axes()    # This creates some axes, so that we
    ax.set_aspect(1) # can set the aspect ratio to 1 i.e. 
                    # x and y axes are scaled equally.                                
    ax.set_xlabel("x coordinate (m)") # Must label axes (with
    ax.set_ylabel("y coordinate (m)") # units) and give
    ax.set_title("Orbit of moon around earth") # plot title.
    #
    ax.plot(results.y[0],results.y[1]) # Make the plot
    ax.plot(0,0)
    ax.legend(['Moon'])                # and add a key.
    ax.axhline(y=0, color='k', linewidth=0.5)
    ax.axvline(x=0, color='k', linewidth=0.5)
    plt.show()

def part2(velocityFactor, orbits):
    #Part two Diffeq
    def f_part2(t, state, Me, Mm, G):
        xm, ym, vx, vy, xp, yp, vpx, vpy = state #all input values

        xpm = xp-xm
        ypm = yp-ym

        dxmdt = vx #moon diffeqs
        dymdt = vy
        dvxdt = -(Me*G*xm)/((xm**2+ym**2)**(3/2)) 
        dvydt = -(Me*G*ym)/((xm**2+ym**2)**(3/2))

        dxpdt = vpx #probe diffeqs
        dypdt = vpy
        dvpxdt = -((Me*G*xp)/((xp**2+yp**2)**(3/2)))-((Mm*G*xpm)/((xpm**2+ypm**2)**(3/2)))
        dvpydt = -((Me*G*yp)/((xp**2+yp**2)**(3/2)))-((Mm*G*ypm)/((xpm**2+ypm**2)**(3/2)))

        return (dxmdt, dymdt, dvxdt, dvydt, dxpdt, dypdt, dvpxdt, dvpydt)


    # Initial Conditions
    Me = 5.97*(10**24)
    Mm = 7.35*(10**22)
    G = 6.67*(10**-11)

    t_min = 0
    t_max = 2360620*int(orbits)
    numpoints = 2000*int(orbits)
    t = np.linspace(t_min, t_max, numpoints)

    rm = 384400000
    vm = sqrt((G*Me)/rm)*float(velocityFactor)

    rpm = 10000000
    vpm = sqrt((G*Mm)/rpm)

    xm0 = rm
    ym0 = 0
    vx0 = 0
    vy0 = vm

    xp0 = rpm+rm #as position is relative to the earth
    yp0 = 0      #not to the moon
    vpx0 = 0
    vpy0 = vm+vpm

    rtol = 1e-6
    atol = 1e-9

    #Solver
    results = solve_ivp(f_part2, (t_min,t_max), (xm0, ym0, vx0, vy0, xp0, yp0, vpx0, vpy0), args=(Me, Mm, G), t_eval=t, atol=atol, rtol=rtol)


    #Graph plotting
    ax=plt.axes()    # This creates some axes, so that we
    ax.set_aspect(1) # can set the aspect ratio to 1 i.e. 
                    # x and y axes are scaled equally.                                
    ax.set_xlabel("x coordinate (m)") # Must label axes (with
    ax.set_ylabel("y coordinate (m)") # units) and give
    ax.set_title("Orbit of moon around earth") # plot title.
    ax.plot(results.y[0],results.y[1], label="Moon" ) # Make the plot
    ax.plot(results.y[4],results.y[5], label="Probe")
    ax.legend(loc='upper right')              # and add a key.
    ax.axhline(y=0, color='k', linewidth=0.5)
    ax.axvline(x=0, color='k', linewidth=0.5)
    plt.show()


MyInput = '0'
while MyInput != 'q':
    MyInput = input('Enter a choice, "1", "2" or "q" to quit: ')
    print('You entered the choice: ',MyInput)
    if MyInput == '1':
        print('You have chosen part (1): simulation of a lunar orbit')
        velocityFactor = input(f'Enter the velocity scalar. A value of 1 will result in a circular orbit, and anything else will give an elliptical orbit: ')
        orbits = input(f'Enter the approximate number of orbits desired: ')
        part1(velocityFactor, orbits)
    elif MyInput == '2':
        print('You have chosen part (2): earth-moon-probe system')
        velocityFactor = input(f'Enter the velocity scalar. A value of 1 will result in a circular orbit, and anything else will give an elliptical orbit: ')
        orbits = input(f'Enter the approximate number of orbits desired: ')
        part2(velocityFactor, orbits)
    elif MyInput != 'q':
        print('This is not a valid choice')
print('You have chosen to finish - goodbye.')