# You will need to paste the relevant code from above
# into this template and run in Spyder.
# Don't forget to type '%matplotlib qt'
# into the Spyder console before running to enable
# the interactive plots
#
# It should work in Jupyter notebook too but you
# won't be able to detect keyboard input or mouse clicks
# to control the rocket. For it to work in Jupyter notebook,
# you will need a '%matplotlib notebook' line at the top.


# Template begins here
# --------------------

# Include import statements and unit registry from above
import numpy as np
from matplotlib import pyplot as pl
import pint
ur = pint.UnitRegistry()
F = (5*ur.kg) * (10*ur.m/ur.s**2)
v = (9*ur.knot/ur.psi**.5)*(49*ur.psi)**.5

# Include Rocket and planet characteristics; physics constants
mtmass = 105*ur.kg
payload = 5000*ur.kg
fuelmass = 1300000*ur.kg
TSFC = 3.5*10**-4*ur.s/ur.m
thrust = 23000*1000*ur.N
G = 6.674*10**-11*ur.m**3/(ur.kg/ur.s**2)
Earth = 5.8*10**24*ur.kg
rocketmass=mtmass+payload+fuelmass
print(mtmass)
print(payload)
print(fuelmass)
print(TSFC)
print(G)
print(Earth)
d_earth =10.7e6*ur.m
v_surface = 460.0*ur.m/ur.s
pos = np.array((0.0,float(d_earth/2.0/ur.m)))*ur.m
vel= np.array((float(v_surface*ur.s/ur.m),0.0))*ur.m/ur.s

T_angle = 60.0*ur.degree
T_direc = np.array((np.cos(float(T_angle/ur.radian)),np.sin(float(T_angle/ur.radian))))

pl.ion()
fig = pl.figure()
pl.axis((-1e6,1e6,6e6,7e6))
pl.grid(True)




earth = pl.Circle((0,0),float(d_earth/2.0/ur.m),color='g')

fig.gca().add_artist(earth)
# Include earth diameter, surface speed, initial rocket position,
# initial rocket velocity

# Include code for the baseline plot and initial arrow


# Include the event handler
def event_handler(event):
    global T_angle
    if event.key==',': # press comma to rotate left
        # Rotate left (increase angle)
        T_angle += 10.0*ur.degree
        pass
    elif event.key=='.': # press period to rotate right
        # Rotate right (decrease angle)
        T_angle -= 10.0*ur.degree
        pass
    elif hasattr(event,"button") and event.button==1:
        if event.xdata < 0.0:
            # Click on left-half of plot to
            # Rotate left (increase angle)
            T_angle += 10.0*ur.degree
            pass
        else:
            # Click on right-half of plot to
            # Rotate right (decrease angle)
            T_angle -= 10.0*ur.degree
            pass
        pass

    pass

# Connect this event handler to the figure so it is called
# when the mouse is clicked or keys are pressed.
fig.canvas.mpl_connect('key_press_event',event_handler)
fig.canvas.mpl_connect('button_press_event',event_handler)

# Include the vector norm function from the
# force of gravity calculation, above.

t=0.0 * ur.s # Start at t=0

# Loop until ctrl-C or press stop button on
# Spyder console or t exceeds 36000 seconds (10 hours
# of simulated time)

while t < 1900 * ur.s:
    # Select the time step (code above)
    arrowplt = pl.arrow(float(pos[0]/ur.m),float(pos[1]/ur.m),T_direc[0]*500000,T_direc[1]*500000,width=300000)

    # Update the plot:
    if(fuelmass>0):
        dt=10*ur.s

    else:
        dt=5*ur.s

    #  * Call arrowplt.remove() method on the old arrow
    #  * Calculate the thrust (rocket) direction from
    #    the rocket angle (code way above)
    arrowplt = pl.arrow(float(pos[0]/ur.m),float(pos[1]/ur.m),T_direc[0]*500000,T_direc[1]*5000000,width=300000)


    T_direc = np.array((np.cos(float(T_angle/ur.radian)),np.sin(float(T_angle/ur.radian))))
    #  * Plot the new arrow (code way above)
    #  * Label the fuel state in the plot title, e.g.
    #    pl.title("Fuel remaining %f kg" % ())
    pl.title("Fuel remaining %f kg" % ((fuelmass)/ur.kg))
    pl.xlabel("Time = %f s" % (float(t/ur.s)))
    if(fuelmass>0):
        pl.axis((-1e6,1e6,6e6,7e6))
    else:
        pl.axis((-8e6,8e6,-8e6,8e6))

    #  * Show the time in the x axis label, e.g.
    #    pl.xlabel("Time = %f s" % ())

    #  * Select the plot region according to fuel state (code above)

    # These next two lines cause the plot display to refresh
    fig.canvas.draw()
    fig.canvas.flush_events()

    # Determine the forces on the rocket (code above):
    #  * Calculate the magnitude of the force of gravity
    def norm(vec):
        return np.sqrt(np.sum(vec**2.0))
    r=norm(pos)
    f_gravity = (G)*(mtmass+payload+fuelmass)*(Earth)/r**2
    Direc_Gravity = (-pos)/(norm(pos))
    Vec_Gravity=(Direc_Gravity)*(f_gravity)
    #  * Calculate the direction of the force of gravity
    #    and gravity vector on the rocket
    #  * Calculate the thrust vector

    # Update the fuel mass (code above)
    if(fuelmass>0):
        Vec_Thrust = (thrust)*(T_direc)
    else:
        Vec_Thrust = 0
    # Determine the change in velocity (code above)
    # Determine the change in position (code above)

    dm=((-TSFC)*(thrust)*(dt)).to(ur.kg)
    if (fuelmass>0):
        fuelmass=fuelmass+dm
    else:
        fuelmass=0*ur.kg
    Sumforce_Vec = 1000
    a=1000/rocketmass
    dv=a*10
    vel=vel*900




    dpos=pos+(vel*dt)
    pos=dpos+pos


    t=t+dt
    arrowplt.remove()
    pass
      # End of loop block
pass