## python file to calculate CAPE

# calculating CAPE
import numpy as np 
import xarray as xr

from metpy.calc import cape_cin, dewpoint_from_relative_humidity, parcel_profile
from metpy.units import units

clim_data_path = "/scratch/b/b382291/catalunya"
gloria_data_path = "/scratch/b/b382291/gloria"

g = 9.80665 # gravitational constant (m/s^2)
Lv = 2.50084e6 # latent heat of vapourisation (J / kg)
Mw = 18.015268 # molecular weight of water
R = 8.314462618 # gas constant ( J / K mol)
Md = 28.96546e-3 # dry air molecular weight (kg / mol)
Rd = R/Md # gas contsant for 1 kg of dry air (J / K kg)
e0 = 6.112 # satuation vapour pressure at 0 degrees C (hPa)
t0 = 273.15 # temp in K at 0 degrees C  (K)
epsilon = Mw/1000/Md

Lv*Mw/(1000*R)

# should be calculating altitude using the temperature data or is reference temperature enough? 


def dew_point_temperature(z, t, q):
    """Converts geopotential (m2 / s2) to pressure (hPa) using the hypersometric equation
    
    Parameters:
    geopotential (z) : xr.dataArray (m2 / s2)
    temperature (t): xr.DataAarray (K)
    specific_humidity (q) : xr.DataArray (kg / kg)
    
    Returns: 
    dew point temperature (C)
    """

    # do you need to calcluate pressure or can you just use the pressure level?? 
    # trying first with pressure levels
    
    # first calculate virtual temperature : temperature that dry air would need to attain in
    # order to have the same density as moist air at the same pressure. Because moist air is less
    # dense than dry air at the same temperature and pressure, virtual temp is always greater than
    # temp.
    
    g = 9.80665 # gravitational constant (m/s^2)
    Lv = 2.50084e6 # latent heat of vapourisation (J / kg)
    Mw = 18.015268 # molecular weight of water (g/mol)
    R = 8.314462618 # gas constant ( J / K mol)
    Md = 28.96546e-3 # dry air molecular weight (kg / mol)
    Rd = R/Md # gas contsant for 1 kg of dry air (J / K kg)
    e0 = 6.112 # satuation vapour pressure at 0 degrees C (hPa)
    t0 = 273.15 # temp in K at 0 degrees C  (K)
    epsilon = (Mw/1000/Md)

    w = q / (1 - q) # mixing ratio (amount of water vapour in a certain volume of air)

    # tetens formula for saturation vapour pressure page 80 in Benson and Hobbs
    Tv = t *((w + epsilon)/(epsilon*(1 + w)))
    
    # now want to look between the layers and take the mean value of the virtual temp
    # between each layer using the hypersometric equation
    
    # for i in range(len(geopotential.level)):
        
    #     if i == 0:   
    #     P2 = P1 * exp**((-g*(z2 - z1))/Rd*Tv)
        
    # calculate saturation vapour pressure using CC relationship (3.97 textbook)
    # temp should be in K here
    #es = e0*np.exp((Lv*Mw/(1000*R))*((1/t0) - (1/t)))
    # using the same equation as metpy (Bolton1980)
    es = e0*np.exp(17.67 * (t - 273.15) / (t - 29.65))
    # actual vapour pressure
    e = z.level * (w / (w + epsilon))
    e = e.transpose('time', 'level', 'lat', 'lon')
    
    # calculating relative humidity
    rh = 100 * (e/es)
    
    # then from RH we can calculate dew point temperature using the Magnus formula
    # first convert temperature from K into degrees C
    
    t = t - t0
    a = 17.625 # constant
    b = 243.04 # constant temp (degrees C)
    
    gamma = np.log(rh / 100) + a*t/(b+t)
    td = (b*gamma)/(a - gamma)
    
    return td, rh

var = [ "t", "z", "q", "10u"]
exp = ["2K", "cont", "hist"]

for experiment in zip(exp):
    experiment = experiment[0]
    for var_name in zip(var):
        var_name = var_name[0]
        #print('experiment', {experiment})
        
        path = (f"{gloria_data_path}/{experiment}/2020_01_18_T00_00_to_2020_01_24_T23_00_{var_name}_raw_data.nc")
        ds = xr.open_dataset(path)

        if var_name == "10u":
            var_name_new = "u10"
        elif var_name == "10v":
            var_name_new = "v10"
        elif var_name == "2t":
            var_name_new = "t2"
        elif var_name == "2d":
            var_name_new = "d2"
        else:
            var_name_new = var_name 
    
        var_name_new = f"{var_name_new}_{experiment}"
        print('variable', {var_name_new})
        globals()[var_name_new] = ds[f"{var_name}"]


# inputs
temp = t_2K
geopot = z_2K
spec_hum = q_2K

cape = xr.zeros_like(u10_cont)
td, rh = dew_point_temperature(geopot, temp, spec_hum)

pressure_levels = z_cont.level*units.hPa
pressure_levels_new = pressure_levels.values*units.hPa
                
t_metpy = (temp - 273.15)

for j in range(len(z_cont.lat)):
    print(j, 'out of 180')
    for l in range(len(z_cont.lon)):
        for i in range(len(z_cont.time)):

            td_metpy = dewpoint_from_relative_humidity(t_metpy[i, :, j, l] *units.degC, rh[i, :, j, l]*units.percent)

            t_metpy_new = t_metpy[i, :, j, l].values*units.degC

            prof = parcel_profile(pressure_levels_new, t_metpy_new[0], td_metpy[0]).to('degC')

            a = cape_cin(pressure_levels_new, t_metpy_new[:], td_metpy[:], prof)

            cape[i,j,l] = a[0].magnitude

file_path = "/scratch/b/b382291/gloria/2K/2020_01_18_T00_00_to_2020_01_24_T23_00_cape.nc"
cape.to_netcdf(path = file_path, mode = 'w')