#
# rocket-launch.py
# perform a launch for the water rocket
#
# hvf@hvf.ch 31.05.17 04.06.2017 06.06.2017


import math
import matplotlib.pyplot as plt

############################################################
# Model parameters

# time step
dt = 0.01       # sec

# initial pressure in air chamber (empty volume)
p0 = 300000     # Pa

# total volume of water rocket
Vtot = 1.5e-3   # 1.5 liter

# empty (air) volume of water rocket; rest is filled with water
V0 = 0.5e-3     # empty volume, 0.5 liter
V0 = 0.4e-3     # empty volume, 0.4 liter
V0 = 0.5e-3     # empty volume, 0.5 liter
V0 = 0.6e-3     # empty volume, 0.6 liter
V0 = 0.75e-3     # empty volume, 0.75 liter
V0 = 0.825e-3     # empty volume, 0.825 liter
V0 = 0.90e-3     # empty volume, 0.9 liter
V0 = 1.00e-3     # empty volume, 1.0 liter
V0 = 0.5e-3     # empty volume, 0.5 liter


# cross section of water rocket (cylinder of 10.5 cm diameter)
A  = 0.105**2*math.pi/4    # m2, diam 10.5 cm

# cross section of bottle opening, 2.5 cm diameter
Q  = 0.025**2*math.pi/4    # m2, diam 2.5 cm

# arbitrary throttling factor such that water does not exit too fast
# we aim at a propulsion time of about 0.3 sec for 1 liter of water
Q = Q / 1.7     # throttling, arbitrary throttling factor

# mass of empty bottle
M  = 0.040      # kg, 40 gr empty weight of rocket (bottle)

# water density
dwater = 1000.  # kg/m3

# air density at sea level, ISA
dair = 1.225    # kg/m3

# drag coefficient of bottle 
Cdrag = 0.6

# some variables and initial values
a0 = 0
v0 = 0
s = 0           # [m], height of rocket above ground
text = ""

# plotting arrays
timelist = []
heights = []
velocs = []
accels = []
jetspeed = []


###########################################################
# doTimeStep - performs one time step of moving rocket

def doTimeStep (dt, p0, V0, Vtot, v0, a0, A, Q, M, dwater, dair, Cdrag):
    """ calculate one time step for a water rocket.
    Input parameters:
    dt  time step [s]
    p0  initial gas pressure [Pa]
    V0  initial gas volume [m3]
    Vtot    total volume of rocket [m3] - Vtot-V0 is volume of water/propellant
    v0  initial vertical velocity of rocket [m/s]
    a0  initial acceleration of rocket [m/s2]
    A   cross section of rocket (skin thickness is zero) [m2]
    Q   nozzle cross section [m2]
    M   empty mass of rocket [kg]
    dwater  density of water/propellant [kg/m3]
    dair    density of ambient air [kg/m3]
    Cdrag   drag coefficient of rocket [-]

    Return values:
    p1  final gas pressure [Pa]
    V1  final gas volume [m3]
    v1  final vertical velocity of rocket [m/s]
    a1  final acceleration of rocket [m/s]
    vj  velocity of water exhaust jet [m/s]
    ds1 increase in vertical position [m]
    dt1 actually used time step 
    text    additional values in a printable string

    Intermediate values:
    m   mass of water/propellant, dwater*(Vtot-V0) [kg]
    T   thrust of water jet [N]
    Vdot    volume flow [m3/s]
    mdot    mass flow [kg/s]
    D   air drag [N]
    """

    g = 9.80665     # earth gravitation [m/s2]
    pAtm = 100000   # atmospheric pressure (rounded)
    # the interior pressure, p0, must be bigger than pAtm
    # or no water is propelled out of the rocket

    text = ""

    # Drag of flying rocket
    D = Cdrag/2.0*dair*v0**2*A

    if (V0 >= Vtot or p0 <= pAtm): 
        vj = 0.0    # no more thrust
        T = 0.0
        V1 = V0
        if (V0 >= Vtot): V1 = Vtot; V0 = Vtot
        Vdot = 0.0
        mdot = 0.0
        # mass of water/propellant - there may still be water but no more p
        m = (Vtot-V1)*dwater
        dm = 0.0
        p1 = 0.0
        if (v0 >= 0.0): 
            D = -D      # D points down
        a1 = -g + D/(M+m)
        adrag = D/(M+m)     # accel/decel due to drag
        text = "{0:3.0f} {1:.3f} {2:.3f} {3:.1f} {4:.2f}".format(0, 0, M+m, D, adrag)
    else:
        # calculate vj from Bernoulli, 1/2 *dwater*vj**2 = p0-pAtm
        vj = math.sqrt(2.*(p0-pAtm)/dwater)

        # volume flow, Vdot = Q*vj
        Vdot = Q*vj

        # mass flow mdot = Q*vj*dwater
        mdot = Vdot*dwater

        # thrust = mdot * vj
        T = mdot * vj

        # check if there is still some water
        # how long does it take (at current conditions) to empty bottle?
        tmax = (Vtot-V0)/Vdot
        if (tmax < dt): dt = tmax   # adjust dt down

        dm = mdot * dt
        dV = Vdot * dt

        # new volume
        V1 = V0 + dV
        # mass of water/propellant
        m = (Vtot-V1)*dwater

        if (v0 >= 0.0): 
            D = -D      # D points down
        adrag = D/(M+m)     # accel/decel due to drag
        # thrust, water loss, total mass, drag
        text = "{0:3.0f} {1:.3f} {2:.3f} {3:.1f} {4:.2f}".format(T, dm, M+m, D, adrag)

        # new pressure, ideal gas
        p1 = p0*V0/V1
    
        # acceleration from forces
        a1 = (T - (M+m)*g - D)/(M+m)

    # intergrate a for v and v for s - may be improved later
    v1 = v0 + a1*dt
    ds = (v0+v1)*dt/2.

    return (p1, V1, v1, a1, vj, ds, dt, text)
###########################################################

# doPrintStep - print lots of stuff for each step

def doPrintStep (t, p, V, v, a, ds, vj, dt, text):
    print ("{0:5.3f} {1:8.1f} {2:7.2e} {3:7.2f}   {4:7.2f} {5:7.2e} {6:7.3e} {7:.3f} {8:s}".format(t, p, V, v, a, ds, vj, dt, text))
        
print ()
print ("  Vtot    V0      A        Q         weight  dwater   dair   Cdrag")
print ("  [m3]    [m3]    [m2]     [m2]      [kg]    [kg/m3]  [kg/m3]      ")

print ("{7:7.2e} {0:7.2e} {1:7.2e} {2:7.2e} {3:7.3f}   {4:5.1f}  {5:7.4f} {6:4.2f}".format(V0, A, Q, M, dwater, dair, Cdrag, Vtot))

print ()
print (" t       p         V         v       a        ds    vj       dt  Thrust mdot  m   Drag brake")
print (" [s]     [Pa]      [m3]      [m/s]   [m/s2]   [m]   [m/s]    [s]  [N]   [kg]  [kg]  [N]   [m/s2]")
p=p0; V=V0; a=a0; v=v0; t=0.0; vj=0.0
doPrintStep (t, p, V, v, a, s, vj, dt, text)

timelist.append(t)
heights.append(s)
velocs.append(v0)
accels.append(a0)
jetspeed.append(0)

# find out maximal values and keep
hmax = 0; thmax = 0
vmax = 0; tvmax = 0
amax = 0; tamax = 0
jmax = 0; tjmax = 0

# let the rocket fly - max 1000 time steps 
for i in range(1000):
    p, V, v, a, vj, ds, dt1, text = doTimeStep (dt, p, V, Vtot, v, a, A, Q, M, dwater, dair, Cdrag)
    t += dt1
    s += ds
    doPrintStep (t, p, V, v, a, s, vj, dt1, text)
    timelist.append(t)
    heights.append(s)
    velocs.append(v)
    accels.append(a)
    jetspeed.append(vj)
    # find maximal values and times
    if (s > hmax): 
        hmax = s; thmax = t
    if (v > vmax): 
        vmax = v; tvmax = t
    if (a > amax): 
        amax = a; tamax = t
    if (vj > jmax): 
        jmax = vj; tjmax = t
    # end of thrust phase: larger time steps
    if (dt1 < dt or vj == 0.0):
        dt = 0.01    #0.01 sec
    # stop if falling below floor
    if (s < 0.0): break
    # increase time steps if falling
    if (v <= 0.0): dt = 0.1
    # break after 3 secs
    #if (t >= 3.0): break

print ("Maximum altitude: {0:.1f} m at {1:.3f} s".format(hmax, thmax))
print ("Maximum velocity: {0:.1f} m/s at {1:.3f} s".format(vmax, tvmax))
print ("Maximum acceleration: {0:.1f} m/s2 at {1:.3f} s".format(amax, tamax))
print ("Maximum jet speed: {0:.1f} m/s at {1:.3f} s".format(jmax, tjmax))

###########################################################
# make a plot with subplots for therelevant stuff

fig = plt.figure(figsize=(10,10))
fig.subplots_adjust(bottom=0.10, left=0.15, top = 0.95, right=0.95)
plt.subplot(221)

plt.title("Altitude, maximum {0:.1f} m".format(hmax))
plt.plot(timelist, heights)
plt.xlabel("Time [s]")
plt.ylabel("Altitude [m]")

plt.subplot(222)
plt.title("Velocity, maximum {0:.1f} m/s".format(vmax))
plt.plot(timelist, velocs)
plt.xlabel("Time [s]")
plt.ylabel("Velocity [m/s]")

plt.subplot(223)
plt.title("Acceleration")
plt.title("Acceleration, maximum {0:.1f} m/s2".format(amax))
plt.plot(timelist, accels)
plt.xlabel("Time [s]")
plt.ylabel("Acceleration [m/s2]")

plt.subplot(224)
plt.title("Jet Speed")
plt.title("Jet Speed, maximum {0:.1f} m/s".format(jmax))
plt.plot(timelist, jetspeed)
plt.xlabel("Time [s]")
plt.ylabel("Jetspeed [m/s]")

# create results file name with pressure p0 and volume V0
resname = "rocket-results-{0:d}-{1:d}.pdf".format(int(p0), int(V0*1e6))
plt.savefig(resname)
print ("Created PDF file ", resname)
plt.show()