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_ngc4008']
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

model_run1().describe()["user_parameters"]

ds_coarsest1 = model_run1().to_dask()
ds_coarsest2 = model_run2().to_dask()
ds_coarsest3 = model_run3().to_dask() 
ds_coarsest4 = model_run4().to_dask() 
ds_coarsest5 = model_run5().to_dask() 
ds_coarsest6 = model_run6().to_dask() 
# ds_fine = model_run(zoom=10,time='PT3H').to_dask()

print('***available time period (in UTC)***')
print('icon_ngc4008',ds_coarsest1.time.min().values,ds_coarsest1.time.max().values)
print('ifs_tco3999-ng5_deepoff',ds_coarsest2.time.min().values,ds_coarsest2.time.max().values)
print('nicam_gl11',ds_coarsest3.time.min().values,ds_coarsest3.time.max().values)
print('mpas_dyamond2',ds_coarsest4.time.min().values,ds_coarsest4.time.max().values)
print('mpas_dyamond3',ds_coarsest5.time.min().values,ds_coarsest5.time.max().values)
print('scream_ne120',ds_coarsest6.time.min().values,ds_coarsest6.time.max().values)

var_longname_dict1 = {
    var: ds_coarsest1[var].attrs.get("long_name", "N/A")
    for var in ds_coarsest1.data_vars
}
var_longname_dict2 = {
    var: ds_coarsest2[var].attrs.get("name", "N/A")
    for var in ds_coarsest2.data_vars
}
var_longname_dict3 = {
    var: ds_coarsest3[var].attrs.get("long_name", "N/A")
    for var in ds_coarsest3.data_vars
}
var_longname_dict4 = {
    var: ds_coarsest4[var].attrs.get("long_name", "N/A")
    for var in ds_coarsest4.data_vars
}
var_longname_dict5 = {
    var: ds_coarsest5[var].attrs.get("long_name", "N/A")
    for var in ds_coarsest5.data_vars
}
var_longname_dict6 = {
    var: ds_coarsest6[var].attrs.get("long_name", "N/A")
    for var in ds_coarsest6.data_vars
}

# MPAS-DYAMOND1/2 (model5) doesn't have latent heat flux
le_vars=['hfls','slhf','hflsd','lh_tavg','hflsd']

ds_fine1 = model_run1(zoom=10,time='PT3H').to_dask()
ds_fine2 = model_run2(zoom=7,time='PT1H').to_dask()
ds_fine3 = model_run3(zoom=9,time='PT3H').to_dask() #no zoomed-in data at PT6H
ds_fine4 = model_run4(zoom=10).to_dask() #hourly
ds_fine6 = model_run6(zoom=7).to_dask() 

#ds_fine2 = model_run2(zoom=11,time='PT1H').to_dask()
ds_fine2=ds_fine2.rename_dims({'value': 'cell'}) #IFS models require renaming for healpix to work

%%time
uxds_fine1 = ux.UxDataset.from_healpix(ds_fine1)
uxds_fine2 = ux.UxDataset.from_healpix(ds_fine2)
uxds_fine3 = ux.UxDataset.from_healpix(ds_fine3)
uxds_fine4 = ux.UxDataset.from_healpix(ds_fine4)
uxds_fine6 = ux.UxDataset.from_healpix(ds_fine6)

#uxda_fine1 #W/m2
#uxda_fine2 #J/m2 1hour -> 3600S ->1/3600W/m2
#uxda_fine3 #W/m2
#uxda_fine4 #W/m2
#uxda_fine6 #s^-3 kg

timestep=18
time_utc=pd.to_datetime(str(uxda_fine2.time.values[timestep]))
if time_utc.tzinfo is None:
    time_utc = time_utc.tz_localize('UTC')

time_mdt = time_utc.tz_convert('US/Mountain')
time_mdt=time_mdt.strftime('%Y.%m.%d, %H:%M:%S')
(uxds_fine2['2t']).isel(time=timestep).plot.points(
    cmap="inferno",
    rasterize=True,
    dynamic=False,
    title=f"model=ifs_tco3999-ng5_deepoff, time={time_mdt} MDT",
    features=["coastline", "borders"]
)

#ifs model
Rnet=uxds_fine2.ssr+uxds_fine2.str # J/m2
H=uxds_fine2.sshf #real Surface sensible heat flux, J/m2

land_mask=ux.UxDataArray(uxds_fine2.sst.isel(time=0).isnull().astype(int),uxgrid=uxds_fine2.uxgrid)

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

sel_le_DJF=uxds_fine2.slhf.sel(time=(uxds_fine2.time.dt.month.isin([1,2,12]))&(uxds_fine2.time.dt.year == 2020)).where(land_mask)/3600 #convert back to W/m2
sel_tp_DJF=uxds_fine2.tp.sel(time=(uxds_fine2.time.dt.month.isin([1,2,12]))&(uxds_fine2.time.dt.year == 2020)).where(land_mask)

sel_le_JJA=uxds_fine2.slhf.sel(time=(uxds_fine2.time.dt.month.isin([6,7,8])&(uxds_fine2.time.dt.year.isin([2020])))).where(land_mask)/3600 #convert back to W/m2
sel_tp_JJA=uxds_fine2.tp.sel(time=(uxds_fine2.time.dt.month.isin([6,7,8])&(uxds_fine2.time.dt.year.isin([2020])))).where(land_mask)

daily_sel_le_DJF=sel_le_DJF.resample(time='1D').mean()
daily_sel_tp_DJF=sel_tp_DJF.resample(time='1D').mean()

daily_sel_le_JJA=sel_le_JJA.resample(time='1D').mean()
daily_sel_tp_JJA=sel_tp_JJA.resample(time='1D').mean()

daily_sel_le_DJF.shape,daily_sel_le_JJA.shape

from scipy.stats import rankdata

ts1_array = daily_sel_le_JJA.values
ts2_array = daily_sel_tp_JJA.values

n_face = ts1_array.shape[1]
corrs_JJA = np.full(n_face, np.nan)

for i in range(n_face):
    ts1 = ts1_array[:, i]
    ts2 = ts2_array[:, i]

    valid = ~np.isnan(ts1) & ~np.isnan(ts2)
    if np.sum(valid) > 2:
        r1 = rankdata(ts1[valid])
        r2 = rankdata(ts2[valid])
        corrs_JJA[i] = np.corrcoef(r1, r2)[0, 1]

np.nanmin(corrs_JJA),np.nanmean(corrs_JJA),np.nanmax(corrs_JJA)

corrs_JJA_ux=ux.UxDataArray(corrs_JJA.astype(float),uxgrid=uxds_fine2.uxgrid,dims=land_mask.dims)

ts1_array = daily_sel_le_DJF.values
ts2_array = daily_sel_tp_DJF.values

n_face = ts1_array.shape[1]
corrs_DJF = np.full(n_face, np.nan)

for i in range(n_face):
    ts1 = ts1_array[:, i]
    ts2 = ts2_array[:, i]

    valid = ~np.isnan(ts1) & ~np.isnan(ts2)
    if np.sum(valid) > 2:
        r1 = rankdata(ts1[valid])
        r2 = rankdata(ts2[valid])
        corrs_DJF[i] = np.corrcoef(r1, r2)[0, 1]

np.nanmin(corrs_DJF),np.nanmean(corrs_DJF),np.nanmax(corrs_DJF)

corrs_DJF_ux=ux.UxDataArray(corrs_DJF.astype(float),uxgrid=uxds_fine2.uxgrid,dims=land_mask.dims)

corrs_DJF_ux.plot()

corrs_JJA_ux.plot()

corrs_ux.plot()