#! /usr/bin/env python3
# -*- coding:utf8 -*-

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

plt.style.use('classic')
plt.rcParams['font.sans-serif'] = 'DejaVu Sans'
plt.rcParams['mathtext.default'] = 'regular'
plt.rcParams['lines.linewidth'] = 2.4

m_bottle = 0.08 # kg
g = 9.81 # m/s**2
V_bottle = 1e-3 # m**3
rho_water = 1e3 # density of water
rho_air = 1.2e0 # density of air
p_atm = 1e5 # one bar in pascal
gamma_air = 1.4 # https://en.wikipedia.org/wiki/Heat_capacity_ratio
A_rocket = 5e-3 # m**2
A_nozzle = 3.8e-4 # m**2
c_drag = 0.35 # https://doi.org/10.1088/0143-0807/30/5/012

def trajectory(frac_water, tt):
    # numerical calculation of a vertical water rocket trajectory
    # considers the mass of air.
    V0_air = V_bottle * (1. - frac_water)
    m0_water = V_bottle * frac_water * rho_water
    m0_air = V0_air * rho_air * (1. + p0 / p_atm)
    p1 = (p0 + p_atm) * (V0_air / V_bottle) ** gamma_air - p_atm
    
    def dif(t, vec):
        y, vy, m_water, m_air = vec
        m_water = max(0., m_water)
        V_water = m_water / rho_water
        V_air = V_bottle - V_water
        
        v_jet = 0.
        mass_rate = 0.
        mass_rate_air = 0.
        if V_air > 0.:
            if m_water > 0.:
                # water expulsion. adiabatic pressure change of air
                p = (p0 + p_atm) * (V0_air / V_air) ** gamma_air - p_atm
                if p >= 0.:
                    v_jet = sqrt(2. * p / rho_water) # Bernoulli's equation
                    mass_rate = rho_water * A_nozzle * v_jet
            else:
                # adiabatic expansion inside the bottle
                p = (p1 + p_atm) * (m_air / m0_air) ** gamma_air - p_atm
                if p > 0.:
                    rho_air_bottle = m_air / V_bottle
                    # air is expanded to atmospheric pressure
                    # use Bernoulli for compressible fluid
                    rho_air_jet = rho_air_bottle * (p_atm/(p_atm+p)) ** (1./gamma_air)
                    v_jet = sqrt(2. / (1 - 1/gamma_air) *
                        ((p+p_atm)/rho_air_bottle - p_atm/rho_air_jet))
                    mass_rate_air = rho_air_jet * A_nozzle * v_jet
        
        m = m_bottle + m_water + m_air - V_bottle * rho_air
        acc = -g
        # add thrust force
        acc += (mass_rate + mass_rate_air) * v_jet / m
        # add drag force of atmosphere
        acc -= 0.5 * rho_air * vy * fabs(vy) * c_drag * A_rocket / m
        
        # introduce damping force close to ground
        y0 = 0.1
        if y < y0:
            omega = sqrt(g / y0)
            acc_damped = -2.*omega*vy - omega**2 * y
            acc = max(acc, acc_damped)
            if y > 0. and vy < 0. and vy**2 >= 2 * y * g:
                acc = max(acc, 0.5 * vy**2 / y)
        
        return [vy, acc, -mass_rate, -mass_rate_air]
    
    # solve the differential equation numerically
    # use a the vector function vec = [y, vy, m_water, m_air]
    sol = solve_ivp(dif, [0., max(tt)], [0., 0., m0_water, m0_air],
        t_eval=tt, method='RK23', atol=1e-7, max_step=0.1)
    # return y(tt), vy(tt), m_water(tt) and m_air(tt)
    return sol.y[0,:], sol.y[1,:], sol.y[2,:], sol.y[3,:]

fig = plt.figure(figsize=(450 / 90.0, 320 / 90.0), dpi=72)
fig.gca().set_position([0.12, 0.15, 0.77, 0.76])
fig.gca().set_prop_cycle(color=['#0072bd', '#d95319', '#edb120', '#7e2f8e'])
m_array = np.linspace(0., 1., 201)

def hmax(mw):
    tarr = np.linspace(0., 10., 41)
    yarr, varr, marr, marr_air = trajectory(mw, tarr)
    
    imax = np.argmax(yarr)
    h_max = yarr[imax] + varr[imax]**2 * 0.5 / g
    return h_max

p0 = 5e5
h_array = [hmax(mw) for mw in m_array]
plt.plot(m_array*100, h_array, label=r'$p_0 = 5$ bar')

p0 = 4e5
h_array = [hmax(mw) for mw in m_array]
plt.plot(m_array*100, h_array, label=r'$p_0 = 4$ bar')

p0 = 3e5
h_array = [hmax(mw) for mw in m_array]
plt.plot(m_array*100, h_array, label=r'$p_0 = 3$ bar')

p0 = 2e5
h_array = [hmax(mw) for mw in m_array]
plt.plot(m_array*100, h_array, label=r'$p_0 = 2$ bar')

plt.grid(True)
plt.xlim(0)
plt.ylim(0)
plt.xlabel('water volume [%]')
plt.ylabel(r'peak altitude [m]')
plt.title('1 liter water rocket')
plt.legend(loc='upper right', handletextpad=0.4, bbox_to_anchor=(1.14, 1.01))
plt.savefig('Water-rocket-altitude-graph.svg')