import numpy as np
import matplotlib.pyplot as plt

# VARIABLES:

v0    = 0.00     # Initial Balloon Velocity in [m/s]
s0    = 0.00     # Initial Balloon Altitude in [m]
p0    = 101325   # (Initial) Air Pressure in [Pa]
rho_a = 1.2250   # (Initial) Air Density in [kg/m^3]

tmax = 50000

T_a   = 288.15   # (Initial) Air Temperature in [K]
g     = 9.80665  # local gravitation in [m/s^2]
dt    = 0.01     # time step in [s]
cD    = 0.47     # drag coefficient balloon (spherical) [-]

R_a = 287.1      # Specific Gas Constant Dry Air in [J/(kg * K)]
R_He = 2077.1    # Specific Gas Constant Helium in [J/(kg * K)]



m_B = 50         # Balloon (Start) Mass in [kg]
m_PL = 10        # Balloon Payload (Start) Mass in [kg]




pi = np.pi

def T(h):
    if h >= 0 and h <= 11000:
        res = 288.15 - 0.0065 * h
        return res
    elif h > 11000 and h <= 20000:
        res = 216.65
        return res
    elif h >= 20000:
        res = 216.65 + 0.0010 * (h - 20000)
        return res

def p(h):
    if h >= 0 and h <= 11000:
        res = 101325 * ((288.15 - 0.0065 * h)/288.15) ** 5.25577
        return res
    elif h > 11000 and h <= 20000:
        res = 22632 * np.exp(-(h - 11000)/6341.62)
        return res
    elif h > 20000:
        res = 5474.87 * ((216.65 + 0.0010 * (z - 20000))/216.65) ** (-34.163)
        return res

def rho(h):
    res = p(h)/(R_a * T(h))
    return res




# p_a = p0         # air pressure = constant
p_He = 101325       # initial gas pressure = air pressure
T_g = 288.15        # gas temperature = air temperature = constant


# GAS DENSITY
rho_g = p_He / (R_He * T_g)

V_B = m_B / rho_g
D_B = (6 * V_B / pi) ** (1/3)


v = []
time = []
alt = []


v_z = v0
z = s0
t = 0
D = 0

rho_plot = []
T_plot = []
p_plot = []
D_plot = []

while t < tmax:
    v.append(v_z)
    alt.append(z)
    time.append(t)

    rho_plot.append(rho(z))
    T_plot.append(T(z))
    p_plot.append(p(z))
    D_plot.append(D)

    D = 0.5 * cD * 0.25 * pi * rho(z) * v_z ** 2 * D_B ** 2
    I = g * V_B * (rho(z) - rho_g)
    G = g * m_PL


    dv_z = (I - np.sign(v_z) * D - G) / m_B * dt

    v_z += dv_z * dt
    z += v_z * dt

    t += dt



#plt.plot(time, rho_plot)
#plt.plot(time, T_plot)
###plt.plot(time, v)
###plt.show()
#plt.plot(time, ind)
#plt.plot(time, D_plot)

plt.plot(time, alt)
plt.show()