import numpy as np
from astropy.time import Time

# SIMPLE ATMOSPHERE MODEL

def T_air(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_air(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 * (h - 20000))/216.65) ** (-34.163)
        return res
def rho_air(h):
    res = p_air(h)/(R_air * T_air(h))
    return res


# USER INPUT

start_height = 0
start_utc = '2006-01-16 08:00:00.000'
start_lat = 0  # start latitude in [deg]
start_lon = 0  # start longitude in [deg]


# INITIALISATION

t = 0
h = start_height
lat = start_lat
lon = start_lon
utc = Time(start_utc)
v_z = 0
p_gas = p_air(h)
T_gas = T_air(h)
T_film = T_gas
p_0 = 101325

dT_gas = (Q_ConvInt / (gamma * m_gas * c_v) - (gamma - 1) / gamma * rho_air(h) * g / (rho_gas * R_gas) * RoC) * dt
dT_film = ((Q_Sun + Q_Albedo + Q_IREarth + Q_IRfilm + Q_ConvExt - Q_ConvInt - Q_IRout) / (c_f * m_film)) * dt



dT_gas = (Q_g/(c_pg * m_g) - (g * M_a * m_g * T_g)/(c_pg * T_a * M_g) * dz/dt) * dt


V_b_new = m_gas * R_gas * T_gas_new/p_gas_new  # r.h.s. with updated values!


v_w = 0

T_g1 = T_g
T_a1 = T_a
# Q_g1 = ???
T_e1 = T_e

# T_g2 = T_g1 + dt * (Q_g1/(c_pg * m_g1) - g * M_a * T_g1 * dh/(c_pg * T_a1 * M_g))
# T_e2 = T_e1 + Q_e/(c_env * m_e) * dt

T_gn = T_g + dt * (Q_g/(c_pg * m_g) - g * M_a * T_g * dh/(c_pg * T_a * M_g))
T_en = T_e + Q_e/(c_env * m_e) * dt

s_n = s + dt * v + 0.5 * a * dt ** 2
v_n = v - v_w + dt * a







T_g = T_gn
T_e = T_en
s = s_n
v = v_n


V = m_g * R_g * T_g / (p_g)



dV_b = V_b_new - V_b  # calculating volume increment
dT_gas = Q_ConvInt/(c_v * m_g) * dt + (gamma - 1) * T_g * (dm_g/m_g - dV_b/V_b)




V_b = V_b_new  # updating current volume


dT_film = ((Q_Sun + Q_Albedo + Q_IREarth + Q_IRfilm + Q_ConvExt - Q_ConvInt - Q_IRout) / (c_f * m_film)) * dt









m_balloon = A_balloon * rho_envelope
m_gas = Vol_balloon * (M_gas * p_gas)/(R_c * T_gas)
m_gross = m_balloon + m_payload
m_tot = m_balloon + m_gas + m_payload

GI = F_buoyant - g * (rho_a - rho_g) * Vol_balloon

W = m_gross * g

m_virtual = m_tot + C_virtual * rho_a * Vol_balloon

R_z = GI + T_z - D_z - W

dV_A = (R_z * dt)/m_virtual

F_Drag = 0.5 * C_D * rho * np.abs(V) * V * A_cross

Re = (rho * V * l)/my

C_D = 24/Re

my = my_0 * (T_0 + C)/(T + C) * (T/T_0) ** (3/2)

C_D = 0.23175 - 0.15757 * f + 0.04744 * f ** 2 - 7.0412 * 10 ** (-3) * f ** 3 + 5.1534 * 10 ** (-4) * f ** 4 - 1.4835 * 10 ** (-5) * f ** 5

Vol_balloon = m_g * R_g * T_g/p_g


Vol_ballon_new = m_g * R_g * T_g/p_g  # r.h.s. with updated values!
dVol = Vol_ballon_new - Vol_balloon  # calculating volume increment
Vol_ballon = Vol_ballon_new  # updating current volume


L_z = g * (M_atm * p_atm/(R * T_atm) - M_gas * p_gas/(R * T_gas)) * Vol_balloon



dT_g = Q_ConvInt/(c_v * m_g) * dt + (gamma - 1) * T_g * (dm_g/m_g - dVol/Vol)

dT_e = (Q_sun + Q_albedo + Q_IREarth + Q_IRsky + Q_ConvExt - Q_ConvInt + Q_IRout)/(c_e * m_e)

Pr_g = my_g * c_pg/kappa_g
Pr_a = my_a * c_pa/kappa_a

beta_g = 1/T_g
beta_a = 1/T_a

Gr_g = d_b ** 3 * rho_g ** 2 * g * beta_g * (T_e - T_g)/(my_g ** 2)
Gr_a = d_b ** 3 * rho_a ** 2 * g * beta_a * (T_e - T_a)/(my_a ** 2)



h = h + dt * V + 0.5 * a * dt ** 2
V = V + dt * a