import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime
begin_time = datetime.now()
from models.sun import sun_angles_analytical
from models.sun import sun_angles_astropy
from models.drag import drag, c_d
from models.simple_atmosphere import T_air, p_air, rho_air
from input.user_input import *
from input.natural_constants import *
from astropy.time import Time
import astropy.units as u
from models.thermal import AirMass

# t: simulation time (0 .. t_max) in [s]
# utc = Time(start_utc) + t * u.second  # current time (UTC)

lat = start_lat   # deg
lon = start_lon   # deg

date = Time('2020-12-13 12:00:00.000')

AZ = sun_angles_astropy(45.0, 45.0, 0, date)[0]
ELV = sun_angles_astropy(45.0, 45.0, 0, date)[1]


# INITIALISATION

t = 0
h = start_height
utc = Time(start_utc)
v_z = 0
p_gas = p_air(h)
T_gas = T_air(h)
T_film = T_gas

t_list = []
h_list = []

while t <= t_end and h >= 0:
    t_list.append(t)
    h_list.append(h)

    rho_gas = p_gas/(R_gas * T_gas)  # calculate gas density through ideal(!) gas equation

    V_b = m_gas/rho_gas  # calculate balloon volume from current gas mass and gas density

    if V_b > V_max:
        V_b = V_max
        m_gas = V_max * rho_gas
    else:
        pass

    m_gross = m_pl + m_film
    m_tot = m_pl + m_film + m_gas
    m_virt = m_tot + c_virt * rho_air(h) * V_b

    d_b = (6.0 * V_b / np.pi) ** (1/3)  # calculate diameter of balloon from its volume
    A_top = np.pi/4 * d_b ** 2

    D = drag(c_d, rho_air(h), d_b, v_z)  # calculate drag force
    I = g * V_b * (rho_air(h) - rho_gas)  # calculate gross inflation
    W = g * m_gross  # calculate weight (force)
    F = I - W - D * np.sign(v_z)

    v_zrel = v_z

    AZ, ELV = sun_angles_analytical(lat, lon, utc)

    A_proj = A_top * (0.9125 + 0.0875 * np.cos(np.pi - 2 * np.deg2rad(ELV)))
    A_surf = 4.94 * V_b ** (2/3)

    A_eff = A_surf
    #A_surf1 = 4.94 * V_b

    # AirMass(x, p_0, ELV, h)

    # CALCULATIONS FOR THERMAL MODEL

    if ELV >= -(180 / np.pi * np.arccos(R_E / (R_E + h))):
        tau_atm = 0.5 * (np.exp(-0.65 * AirMass(p_air(h), p_0, ELV, h)) + np.exp(-0.095 * AirMass(p_air(h), p_0, ELV, h)))
        tau_atmIR = 1.716 - 0.5 * (np.exp(-0.65 * p_air(h) / p_0) + np.exp(-0.095 * p_air(h) / p_0))
    else:
        tau_atm = 0
        tau_atmIR = 0

    doy = int(utc.doy)

    MA = (357.52911 + 0.98560028 * (utc.jd - 2451545)) % 360  # in degree, reference: see folder "literature"
    TA = MA + 2 * e * np.sin(np.deg2rad(MA)) + 5 / 4 * e ** 2 * np.sin(np.deg2rad(2 * MA))
    I_Sun = 1367.5 * ((1 + e * np.cos(np.deg2rad(TA))) / (1 - e ** 2)) ** 2
    I_SunZ = I_Sun * tau_atm
    q_sun = I_SunZ
    q_IRground = epsilon_ground * sigma * T_ground ** 4
    q_IREarth = q_IRground * tau_atmIR

    if ELV <= 0:
        q_Albedo = 0
    else:
        q_Albedo = Albedo * I_Sun * np.sin(np.deg2rad(ELV))

    my_air = (1.458 * 10 ** -6 * T_air(h) ** 1.5) / (T_air(h) + 110.4)
    my_gas = 1.895 * 10 ** -5 * (T_gas / 273.15) ** 0.647
    k_air = 0.0241 * (T_air(h) / 273.15) ** 0.9
    k_gas = 0.144 * (T_gas / 273.15) ** 0.7
    Pr_air = 0.804 - 3.25 * 10 ** (-4) * T_air(h)
    Pr_gas = 0.729 - 1.6 * 10 ** (-4) * T_gas

    Gr_air = (rho_air(h) ** 2 * g * np.abs(T_film - T_air(h)) * d_b ** 3) / (T_air(h) * my_air ** 2)
    Nu_air = 2 + 0.45 * (Gr_air * Pr_air) ** 0.25
    HC_free = Nu_air * k_air / d_b
    Re = np.abs(v_zrel) * d_b * rho_air(h) / my_air

    HC_forced = k_air / d_b * (2 + 0.41 * Re ** 0.55)
    HC_internal = 0.13 * k_gas * ((rho_gas ** 2 * g * np.abs(T_film - T_gas) * Pr_gas) / (T_gas * my_air ** 2)) ** (
                1 / 3)
    HC_external = np.maximum(HC_free, HC_forced)

    HalfConeAngle = np.arcsin(R_E / (R_E + h))
    ViewFactor = (1 - np.cos(HalfConeAngle)) / 2

    Q_Sun = alpha_VIS * A_proj * q_sun * (1 + tau_VIS / (1 - r_VIS))
    Q_Albedo = alpha_VIS * A_surf * q_Albedo * ViewFactor * (1 + tau_VIS / (1 - r_VIS))
    Q_IREarth = alpha_IR * A_surf * q_IREarth * ViewFactor * (1 + tau_IR / (1 - r_IR))
    Q_IRfilm = sigma * epsilon * alpha_IR * A_surf * T_film ** 4 * 1 / (1 - r_IR)
    Q_IRout = sigma * epsilon * A_surf * T_film ** 4 * (1 * tau_IR / (1 - r_IR))
    Q_ConvExt = HC_external * A_eff * (T_air(h) - T_film)
    Q_ConvInt = HC_internal * A_eff * (T_film - T_gas)

    RoC = -v_z

    dT_gas = (Q_ConvInt / (gamma * R_gas) - (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

    ###

    dv_z = F/m_virt * dt  # velocity increment

    v_z += dv_z * dt  # velocity differential
    h += v_z * dt  # altitude differential

    p_gas = p_air(h)
    T_gas = T_gas + dT_gas * dt
    T_film = T_film + dT_film * dt

    t += dt  # time increment
    utc = Time(start_utc) + t * u.second

plt.plot(t_list, h_list)
plt.show()