import numpy as np
import math
import xarray as xr
import matplotlib.pyplot as plt
from scipy.spatial import cKDTree
from scipy import interpolate
import cartopy.crs as ccrs
from datetime import datetime
start = datetime.now()
print(start)
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 input.user_input import *
from input.natural_constants import *
from astropy.time import Time
import astropy.units as unit
from models.thermal import AirMass
from netCDF4 import Dataset
import pandas as pd

data = pd.read_excel(r'C:\Users\marcel\PycharmProjects\MasterThesis\Mappe1.xls', sheet_name='Tabelle3')

comp_time = pd.DataFrame(data, columns= ['Time'])
comp_height = pd.DataFrame(data, columns= ['Height'])


def transform(lon, lat, t):
    # WGS 84 reference coordinate system parameters
    A = 6378137.0  # major axis [m]
    E2 = 6.69437999014e-3  # eccentricity squared

    t_s = 1000 * t

    lon_rad = np.radians(lon)
    lat_rad = np.radians(lat)
    # convert to cartesian coordinates
    r_n = A / (np.sqrt(1 - E2 * (np.sin(lat_rad) ** 2)))
    x = r_n * np.cos(lat_rad) * np.cos(lon_rad)
    y = r_n * np.cos(lat_rad) * np.sin(lon_rad)
    z = r_n * (1 - E2) * np.sin(lat_rad)
    return x, y, z, t_s

data = Dataset("test2021.nc", "r", format="NETCDF4")

ERAtime = data.variables['time'][:]      # time
ERAlat = data.variables['latitude'][:]   # latitude
ERAlon = data.variables['longitude'][:]  # longitude
ERAz = data.variables['z'][:]/g          # geopotential to geopotential height
ERApress = data.variables['level'][:]    # pressure level
ERAtemp = data.variables['t'][:]         # temperature in K

vw_x = data.variables['u'][:]  # v_x
vw_y = data.variables['v'][:]  # v_y
vw_z = data.variables['w'][:]  # v_z

lon_era2d, lat_era2d, time_era = np.meshgrid(ERAlon, ERAlat, ERAtime)

xs, ys, zs, ts = transform(lon_era2d.flatten(), lat_era2d.flatten(), time_era.flatten())

tree = cKDTree(np.column_stack((xs, ys, zs, ts)))  # !




lat = start_lat   # deg
lon = start_lon   # deg

t = 0  # simulation time in seconds
h = start_height
utc = Time(start_utc)
epoch_diff = (Time(start_utc).jd - Time('1900-01-01 00:00:00.0').jd) * 24.000000
t_epoch = epoch_diff



xt, yt, zt, tt = transform(lon, lat, t)  # test coordinates

d, inds = tree.query(np.column_stack((xt, yt, zt, tt)), k=8)  # longitude, latitude, time in h  # !
w = 1.0 / d[0] ** 2

lat_sel = np.unravel_index(inds[0], lon_era2d.shape)[0]
lon_sel = np.unravel_index(inds[0], lon_era2d.shape)[1]
time_sel = np.unravel_index(inds[0], lon_era2d.shape)[2]

interp4d_height = np.ma.dot(w, ERAz[time_sel, :, lat_sel, lon_sel]) / np.sum(w)
interp4d_temp = np.ma.dot(w, ERAtemp[time_sel, :, lat_sel, lon_sel]) / np.sum(w)
interp4d_vw_x = np.ma.dot(w, vw_x[time_sel, :, lat_sel, lon_sel]) / np.sum(w)
interp4d_vw_y = np.ma.dot(w, vw_y[time_sel, :, lat_sel, lon_sel]) / np.sum(w)
interp4d_vw_z = np.ma.dot(w, vw_z[time_sel, :, lat_sel, lon_sel]) / np.sum(w)

pressure_hPa = np.array([1, 2, 3, 5, 7, 10, 20, 30, 50, 70, 100, 125, 150, 175, 200,225, 250, 300, 350, 400,
                        450, 500, 550, 600, 650, 700, 750, 775, 800, 825, 850, 875, 900, 925, 950, 975, 1000])

pressure = 100 * pressure_hPa
print(pressure)


# height_interp1d = interpolate.interp1d(interp4d_height, pressure)
temp_interp1d = interpolate.interp1d(interp4d_height, interp4d_temp)
press_interp1d = interpolate.interp1d(interp4d_height, pressure)
vw_x_interp1d = interpolate.interp1d(interp4d_height, interp4d_vw_x)
vw_y_interp1d = interpolate.interp1d(interp4d_height, interp4d_vw_y)
vw_z_interp1d = interpolate.interp1d(interp4d_height, interp4d_vw_z)


#u = 0
#v = 0
#w = 0
T_air = temp_interp1d(h)
p_air = press_interp1d(h)
rho_air = p_air/(R_air * T_air)

u = vw_x_interp1d(h)
v = vw_y_interp1d(h)
w = -1 / g * vw_z_interp1d(h) * R_air * T_air / p_air


v_x = 0
v_y = 0
v_z = 0

v_rel = 0

T_gas = T_air
T_film = T_gas

t_list = []
h_list = []
v_list = []
lat_list = []
lon_list = []
p_list = []
Temp_list = []
rho_list = []

while t <= t_end and h >= 0:
    t_list.append(t)
    h_list.append(h)
    v_list.append(v_rel)
    lat_list.append(lat)
    lon_list.append(lon)

    p_list.append(p_air)
    Temp_list.append(T_air)
    rho_list.append(rho_air)

    t_epoch = epoch_diff + t/3600

    xt, yt, zt, tt = transform(lon, lat, t_epoch)  # current balloon coordinates in cartesian coordinates: x [m], y [m], z [m], time [h since 1900-01-01 00:00:00.0]

    d, inds = tree.query(np.column_stack((xt, yt, zt, tt)), k=8)  # longitude, latitude, time in h  # !
    w = 1.0 / d[0] ** 2

    lat_sel = np.unravel_index(inds[0], lon_era2d.shape)[0]
    lon_sel = np.unravel_index(inds[0], lon_era2d.shape)[1]
    time_sel = np.unravel_index(inds[0], lon_era2d.shape)[2]

    interp4d_height = np.ma.dot(w, ERAz[time_sel, :, lat_sel, lon_sel]) / np.sum(w)
    interp4d_temp = np.ma.dot(w, ERAtemp[time_sel, :, lat_sel, lon_sel]) / np.sum(w)
    interp4d_vw_x = np.ma.dot(w, vw_x[time_sel, :, lat_sel, lon_sel]) / np.sum(w)
    interp4d_vw_y = np.ma.dot(w, vw_y[time_sel, :, lat_sel, lon_sel]) / np.sum(w)
    interp4d_vw_z = np.ma.dot(w, vw_z[time_sel, :, lat_sel, lon_sel]) / np.sum(w)

    press_interp1d = interpolate.interp1d(interp4d_height, pressure)
    temp_interp1d = interpolate.interp1d(interp4d_height, interp4d_temp)
    vw_x_interp1d = interpolate.interp1d(interp4d_height, interp4d_vw_x)
    vw_y_interp1d = interpolate.interp1d(interp4d_height, interp4d_vw_y)
    vw_z_interp1d = interpolate.interp1d(interp4d_height, interp4d_vw_z)


    try:
        p_air = press_interp1d(h)
        T_air = temp_interp1d(h)
    except:
        if (h > interp4d_height[0]):
            h = interp4d_height[0]
        else:
            h = interp4d_height[-1]

        p_air = press_interp1d(h)
        T_air = temp_interp1d(h)

        #print("height")
        #print(h)
        #print("temperature")
        #print(temp_interp1d(h))
        #print("pressure")
        #print(press_interp1d(h))
        #print("vw_x")
        #print(vw_x_interp1d(h))

    p_gas = p_air
    rho_air = p_air / (R_air * T_air)
    u = vw_x_interp1d(h)
    v = vw_y_interp1d(h)
    w = -1/g * vw_z_interp1d(h) * R_air * T_air / p_air

    v_relx = u - v_x
    v_rely = v - v_y
    v_relz = w - v_z

    v_rel = (v_relx ** 2 + v_rely ** 2 + v_relz ** 2) ** (1/2)

    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_design:
       V_b = V_design
       #m_gas = V_design * 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 * V_b

    d_b = 1.383 * V_b ** (1/3)  # calculate diameter of balloon from its volume
    L_goreB = 1.914 * V_b ** (1/3)
    h_b = 0.748 * d_b
    A_surf = 4.94 * V_b ** (2/3)
    A_surf1 = 4.94 * V_design ** (2/3) * (1 - np.cos(np.pi * L_goreB/L_goreDesign))
    A_eff = 0.65 * A_surf + 0.35 * A_surf1
    A_top = np.pi/4 * d_b ** 2

    D = drag(c_d, rho_air, d_b, v_rel)  # calculate drag force

    if v_rel == 0:
        Drag_x = 0
        Drag_y = 0
        Drag_z = 0
    else:
        Drag_x = D * v_relx/v_rel
        Drag_y = D * v_rely/v_rel
        Drag_z = D * v_relz/v_rel

    I = g * V_b * (rho_air - rho_gas)  # calculate gross inflation
    W = g * m_gross  # calculate weight (force)
    F = I - W + Drag_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)))

    # 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, p_0, ELV, h)) + np.exp(-0.095 * AirMass(p_air, p_0, ELV, h)))
        tau_atmIR = 1.716 - 0.5 * (np.exp(-0.65 * p_air / p_0) + np.exp(-0.095 * p_air / 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 ** 1.5) / (T_air + 110.4)
    my_gas = 1.895 * 10 ** -5 * (T_gas / 273.15) ** 0.647
    k_air = 0.0241 * (T_air / 273.15) ** 0.9
    k_gas = 0.144 * (T_gas / 273.15) ** 0.7
    Pr_air = 0.804 - 3.25 * 10 ** (-4) * T_air
    Pr_gas = 0.729 - 1.6 * 10 ** (-4) * T_gas

    Gr_air = (rho_air ** 2 * g * np.abs(T_film - T_air) * d_b ** 3) / (T_air * 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_relz) * d_b * rho_air / 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 - T_film)
    Q_ConvInt = HC_internal * A_eff * (T_film - T_gas)

    # RoC = -v_z
    # 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
    #  v_w = w
    #  v = v_z

    s = h

    a_x = Drag_x/m_virt
    a_y = Drag_y/m_virt
    a_z = F/m_virt

    # DIFFERENTIAL EQUATIONS

    dx = dt * v_x + 0.5 * a_x * dt ** 2   # lon
    dy = dt * v_y + 0.5 * a_y * dt ** 2   # lat

    lon = lon + np.rad2deg(np.arctan(dx/((6371229.0 + h) * np.cos(np.deg2rad(lat)))))
    lat = lat + np.rad2deg(np.arctan(dy/(6371229.0 + h)))

    v_xn = v_x + dt * a_x
    v_yn = v_y + dt * a_y

    s_n = s + dt * v_z + 0.5 * a_z * dt ** 2
    dh = s_n - s
    v_n = v_z + dt * a_z

    T_gn = T_gas + dt * (Q_ConvInt / (gamma * c_v * m_gas) - g * R_gas * T_gas * dh / (c_v * gamma * T_air * R_air))
    T_en = T_film + (Q_Sun + Q_Albedo + Q_IREarth + Q_IRfilm + Q_ConvExt - Q_ConvInt - Q_IRout) / (c_f * m_film) * dt

    T_g = T_gn
    T_e = T_en
    s = s_n
    v_x = v_xn
    v_y = v_yn
    v_z = v_n

    h = s
    T_gas = T_g
    T_film = T_e
#   v_z = v

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



#print(len(lon_list), len(lat_list))
#"""
#plt.plot(start_lon, start_lat, 'rx')
#plt.plot(ax=ax, lon_list, lat_list, transform=ccrs.PlateCarree())
#plt.show()
#"""


### WORKS!
###
###dset = xr.open_dataset("test2021.nc")
###
###print(dset['t'][1][0])  # [time] [level]
###
#fig = plt.figure()  #figsize=[120,50])
###
#ax = fig.add_subplot(111, projection=ccrs.PlateCarree())
###
###dset['t'][1][0].plot(ax=ax, cmap='jet', transform=ccrs.PlateCarree())
###ax.coastlines()
###plt.show()

#print(len(ERAlat))
#print(len(ERAlon))
#print(len(ERAz))





#"""
plt.plot(comp_time, comp_height, 'r--', label='PoGo Flight Test')
#plt.plot(t_list, v_list, 'k--', label='Ballon v_rel')
plt.plot(t_list, h_list, 'k-', label='Balloon Altitude')
#plt.plot(t_list, p_list, 'r-', label='Air Pressure [Pa]')
#plt.plot(t_list, Temp_list, 'b-', label='Air Temperature [K]')
#plt.plot(t_list, rho_list, 'g-', label='Air Density in [kg/m$^3$]')
plt.xlabel('time in s')
plt.ylabel('Balloon Altitude in m')
plt.legend()
plt.show()
#"""