import numpy as np
import matplotlib.pyplot as plt

g = 9.81

# F = W*V_e / g
# I_SP = V_e / g
# V_burnout = 1000 = I_SP*G ln(W_L/W_BO)

def specific_impulse(v_burnout, w_launch, w_burnout):
    return v_burnout / (g*np.log(w_launch/w_burnout))

def v_burnout(I_sp, t_burn, w_launch, w_burnout):
    return I_sp * g*np.log(w_launch/w_burnout)


def exit_velocity():
    plt.figure(figsize=(16,9))
    v_burnout = 1000
    w_rocket = 300
    w_propellant_0 = 800

    w_lauch = w_rocket + w_propellant_0

    I_sp = specific_impulse(v_burnout, w_lauch, w_rocket)
    v_exit = I_sp * g

    print("I_sp:", I_sp)
    print("V_e:", v_exit)
    
    # Assume constant weight flow rate
    a_burnout = 18*g
    weight_flow_rate = a_burnout*w_rocket/v_exit

    dt = 0.1
    t = np.arange(0, 12, dt)

    w_propellant = w_propellant_0 - weight_flow_rate*t
    w_propellant = np.maximum(w_propellant, 0)
    
    thrust = weight_flow_rate*v_exit/g
    w_total = w_rocket + w_propellant
    # Acceleration in g is force/weight
    acceleration_g = thrust/w_total

    v_no_gravity = np.cumsum(acceleration_g*g)
    v_gravity_45_x = v_no_gravity*np.cos(np.pi/4)
    v_gravity_45_y = v_no_gravity*np.sin(np.pi/4) - np.cumsum(g*np.ones_like(acceleration_g))
    v_gravity_45 = np.sqrt(v_gravity_45_x**2 + v_gravity_45_y**2)

    plt.subplot(411)
    plt.plot(t, w_propellant)
    plt.ylabel("Propellant weight [N]")
    plt.subplot(412)
    plt.plot(t, acceleration_g)
    plt.ylabel("Acceleration [g]")
    plt.subplot(413)
    plt.plot(t, v_no_gravity)
    plt.ylabel("Velocity w/o gravity [m/s]")
    plt.subplot(414)
    plt.plot(t, v_gravity_45)
    plt.ylabel(r"Velocity w/ gravity [m/s]")
    plt.xlabel("Time [s]")
    plt.savefig("rocket_motor.png")
    plt.show()

def main():
    I_sp = 250
    t_burn = 14
    t = np.linspace(0, t_burn, 500)
    m_0 = 1200
    m_final = 700
    
    m_propellant = m_0 - m_final
    w_propellant_0 = g*m_propellant
    w_rocket = g*m_final

    weight_flow_rate = w_propellant_0/t_burn
    v_exit = I_sp*g
    w_propellant = w_propellant_0 - weight_flow_rate*t

    thrust = weight_flow_rate*v_exit*g
    w_total = w_rocket + w_propellant
    
    acceleration_g = thrust/w_total
    
    plt.plot(t, acceleration_g)
    plt.savefig("timed_burn.png")
    plt.show()

    


if __name__ == "__main__":
    main()