import datetime
current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
print(f'Last updated at {current_time}')

from pathlib import Path 
import xarray as xr
import cartopy.crs as ccrs
import uxarray as ux
import numpy as np
import cartopy.feature as cf
import matplotlib.pyplot as plt
import pandas as pd

import intake

cat_url = "https://digital-earths-global-hackathon.github.io/catalog/catalog.yaml"
cat = intake.open_catalog(cat_url).NCAR

# check on global models
model_run1 = cat['icon_d3hp003']
model_run2 = cat['ifs_tco3999-ng5_deepoff']
model_run3 = cat['nicam_gl11']
model_run4 = cat['mpas_dyamond3']
model_run5 = cat['mpas_dyamond2']
model_run6 = cat['scream_ne120'] #3hourly average

ds_fine1 = model_run1(zoom=7,time='P1D').to_dask()

uxds_fine1 = ux.UxDataset.from_healpix(ds_fine1)

var_longname_dict1 = {
    var: ds_fine1[var].attrs.get("long_name", "N/A")
    for var in ds_fine1.data_vars
}

var_longname_dict1

# 'pres_sfc': 'surface pressure'
#  'rlds': 'surface downwelling longwave radiation',
#  'rlus': 'surface upwelling longwave radiation'
# 'rsds': 'surface downwelling shortwave radiation',
#  'rsus': 'surface upwelling shortwave radiation',
# 'hfls': 'latent heat flux',
#  'hfss': 'sensible heat flux',
# 'tas': 'temperature in 2m',

land_mask=ux.UxDataArray((uxds_fine1.sftlf==1).astype(int),uxgrid=uxds_fine1.uxgrid)

land_mask.plot.points(
    cmap="inferno",
    rasterize=True,
    dynamic=False,
    features=["coastline", "borders"]
)

uxds_fine1=uxds_fine1.sel(time=uxds_fine1.time.dt.year==2020)

%%time
sel_vswc=(uxds_fine1['mrso'])
sel_le=(uxds_fine1['hflsd']) #W/m2
sel_h=(uxds_fine1['hfssd']) #W/m2
sel_rnet=((uxds_fine1['rlds']-uxds_fine1['rlus']+uxds_fine1['rsds']-uxds_fine1['rsus'])) #W/m2
sel_t2m=(uxds_fine1['tas']-273.15) #C
sel_sp=(uxds_fine1['ps']/1000) #kPa

%%time
sel_le_DJF=sel_le.sel(time=(uxds_fine1.time.dt.month.isin([1,2,12]))&(uxds_fine1.time.dt.year == 2020)).where(land_mask)
sel_h_DJF=sel_h.sel(time=(uxds_fine1.time.dt.month.isin([1,2,12]))&(uxds_fine1.time.dt.year == 2020)).where(land_mask)
sel_rnet_DJF=sel_rnet.sel(time=(uxds_fine1.time.dt.month.isin([1,2,12]))&(uxds_fine1.time.dt.year == 2020)).where(land_mask)
sel_t2m_DJF=sel_t2m.sel(time=(uxds_fine1.time.dt.month.isin([1,2,12]))&(uxds_fine1.time.dt.year == 2020)).where(land_mask)
sel_sp_DJF=sel_sp.sel(time=(uxds_fine1.time.dt.month.isin([1,2,12]))&(uxds_fine1.time.dt.year == 2020)).where(land_mask)

%%time
sel_le_JJA=sel_le.sel(time=(uxds_fine1.time.dt.month.isin([6,7,8])&(uxds_fine1.time.dt.year.isin([2020])))).where(land_mask)
sel_h_JJA=sel_h.sel(time=(uxds_fine1.time.dt.month.isin([6,7,8])&(uxds_fine1.time.dt.year.isin([2020])))).where(land_mask)
sel_rnet_JJA=sel_rnet.sel(time=(uxds_fine1.time.dt.month.isin([6,7,8])&(uxds_fine1.time.dt.year.isin([2020])))).where(land_mask)
sel_t2m_JJA=sel_t2m.sel(time=(uxds_fine1.time.dt.month.isin([6,7,8])&(uxds_fine1.time.dt.year.isin([2020])))).where(land_mask)
sel_sp_JJA=sel_sp.sel(time=(uxds_fine1.time.dt.month.isin([6,7,8])&(uxds_fine1.time.dt.year.isin([2020])))).where(land_mask)

def saturation_vapor_pressure(T):
    return 6.112*np.exp((17.67*T)/(T+243.5))*0.1  # hPa -> kPa

def delta_svp(T):
    es=saturation_vapor_pressure(T)
    return (4302.645*es)/((T+243.5)**2)  # in kPa/°C

def alpha_piecewise(vswc): # values from literature based on Priestley–Taylor equation
    return xr.where(
        vswc < 0.1, 0.8,
        xr.where(vswc < 0.2, 1.0, 1.26)
    )

def calc_gamma(P_kPa):
    cp=1.013e3      # J/kg/°C
    epsilon=0.622
    lambda_v=2.45e6 # Latent heat of vaporization J/kg
    return (cp*P_kPa)/(epsilon*lambda_v)  # in kPa/°C

def priestley_taylor_ET(vswc, Rnet, T, P_kPa): # assume Rnet=LE+H
    alpha=alpha_piecewise(vswc)
    delta=delta_svp(T)
    gamma=calc_gamma(P_kPa)
    coeff=alpha*(delta/(delta+gamma))
    return coeff*Rnet  # W/m2

def safe_corr(x, y, dim='time', min_count=30):
    valid = x.notnull() & y.notnull()
    count = valid.sum(dim=dim)
    r = xr.corr(x, y, dim=dim)
    return r.where(count >= min_count)
    
def calculate_coupling_index(vswc, Rnet, LE, T, P):
    H=Rnet-LE

    LE_p=priestley_taylor_ET(vswc, Rnet, T, P)
    H_p=Rnet-LE_p
    
    H_xr=H.to_dataset(name='H')
    Hp_xr=H_p.to_dataset(name='Hp')
    T_xr=T.to_dataset(name='T2M')

    H_corr = safe_corr(H, T, dim='time', min_count=30)
    Hp_corr = safe_corr(H_p, T, dim='time', min_count=30)

    PI=H_corr-Hp_corr
    PI.name='coupling_index'

    return PI

%%time
PI_DJF=calculate_coupling_index(sel_rnet_DJF,sel_le_DJF,sel_t2m_DJF,sel_sp_DJF).compute()

PI_JJA=calculate_coupling_index(sel_rnet_JJA,sel_le_JJA,sel_t2m_JJA,sel_sp_JJA).compute()

PI_DJF.plot()

PI_JJA.plot()