#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Nov 17 14:37:43 EST 2023
if __name__ == '__main__':
    import sys,os
    from misc.timer import Timer
    tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
#from misc import get_kws_from_argv
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#vapor pressure deficit (from Chiodi et al. 2021  https://doi.org/10.1029/2021GL092830)
def func_vpd(t_ref, rh_ref):
    """vapor pressure deficit (from Chiodi et al. 2021  https://doi.org/10.1029/2021GL092830)
    es = 6.109*np.exp(17.625*T/(T+243.04)) #units: hPa; T units is degC (not K!)
    vpd = es*(1 - RH/100) #units hPa; RH units is %"""
    #t_ref = ds['t_ref'].load() #units degC
    #rh_ref = ds['rh_ref'].load() #units %
    RH = rh_ref
    T = t_ref - 273.15 #K -> C
    es = 6.109*np.exp(17.625*T/(T+243.04)) #units: hPa
    vpd = es*(1-RH/100)
    vpd.attrs['units'] = 'hPa'
    return vpd
ofile = 'data_AM2.5C360_amipHadISSTrcp45_tigercpu_intelmpi_18_1080PE_3ens_1871-2100_atmos_daily.region.199_206_18_23.vpd.nc'
if os.path.exists(ofile): print('[exists]:', ofile); sys.exit()
#t_ref_max
ifile = 't_ref_max_AM2.5C360_amipHadISSTrcp45_tigercpu_intelmpi_18_1080PE_3ens_1871-2100_atmos_daily.region.199_206_18_23.nc'
da = xr.open_dataarray(ifile).load()
t_ref_max = da
#rh_ref
ifile = 'rh_ref_AM2.5C360_amipHadISSTrcp45_tigercpu_intelmpi_18_1080PE_3ens_1871-2100_atmos_daily.region.199_206_18_23.nc'
da = xr.open_dataarray(ifile).load()
rh_ref = da
#wind
ifile = 'wind_AM2.5C360_amipHadISSTrcp45_tigercpu_intelmpi_18_1080PE_3ens_1871-2100_atmos_daily.region.199_206_18_23.nc'
da = xr.open_dataarray(ifile).load()
wind = da

#vpd calculation
da = func_vpd(t_ref_max, rh_ref)
vpd = da
#save
da.to_dataset(name='vpd').to_netcdf(ofile)
print('[saved]:', ofile)

#hdw
daname = 'hdwi'
hdwi = wind*vpd
hdwi.attrs['units'] = 'hPa m/s'
#save
ofile_hdwi = ofile.replace('vpd', daname)
hdwi.to_dataset(name=daname).to_netcdf(ofile_hdwi)
print('[saved]:', ofile_hdwi)
 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()