import math

#rho = Air density
#dt = Time step
def simulate_rocket(thrust, burn_time, mass_dry, mass_propellant, Cd, A, rho=1.225, dt=0.01):
    
    # Initial conditions
    t = 0.0 #Time
    v = 0.0 #Velocity
    h = 0.0 #Height
    g = 9.81 #Gravity in m/s
    mass = mass_dry + mass_propellant #(kg)
    mass_flow_rate = mass_propellant / burn_time  #(kg/s)

    # Start simulating with the params.
    while True:
        if t < burn_time: #If we still need to burn fuel, burn some.
            #Also remember to subtract from the mass
            current_thrust = thrust
            mass -= mass_flow_rate * dt
        else:
            current_thrust = 0 #No thrust added, no fuel burning.

        # Drag force calculations
        drag = 0.5 * rho * Cd * A * v**2
        if v > 0:
            drag *=1
        else:
            drag *=-1

        # Acceleration (upwards positive)
        a = (current_thrust - drag - mass * g) / mass
        v += a * dt
        h += v * dt
        t += dt

        # Stop simulation where velocity starts to become negative
        if v <= 0 and t > burn_time:
            break

    return h, t

#Propulsion + Avionics + Cables estimate + Nose Cone + Avionics Structure + External structure + Recovery estimate
total_rocket_weight = 0.210 +0.015 + 0.050 + 0.0115 +0.048 + 0.100

apogee, flight_time = simulate_rocket(
    thrust=36, # N
    burn_time=2.1, # s
    mass_dry=total_rocket_weight, # kg
    mass_propellant=0.06, # kg
    Cd=0.75, # dimensionless
    A=0.00212372 # m²
)

print(f"Estimated dry weight: {total_rocket_weight} kg")
print(f"Estimated apogee: {apogee:.5f} m")
print(f"Total flight time until apogee: {flight_time:.5f} s")