#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Nov 18 17:32:17 EST 2024
if __name__ == '__main__':
    import sys,os
    try:
        from misc.timer import Timer
        tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
    except:
        pass
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__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
daname_partial = get_kws_from_argv('daname_partial', 'PI')
ofile = f'prob{daname_partial}_HIRAM_amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE_ens01_0850-1999.nc'
print('ofile:', ofile)
if os.path.exists(ofile) and 'od' not in sys.argv:
    print('[exists]:', ofile)
    sys.exit()
#PI
ifile = 'PI_HIRAM_amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE_ens01_0850-1999.nc'
print('ifile:', ifile)
print('PI loading...')
da = xr.open_dataarray(ifile)#.load()
PI = da
#Vshear
ifile = 'Vshear_HIRAM_amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE_ens01_0850-1999.nc'
print('ifile:', ifile)
print('Vshear loading...')
da = xr.open_dataarray(ifile)#.load()
Vshear = da
#chi
ifile = 'chi_HIRAM_amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE_ens01_0850-1999.nc'
print('ifile:', ifile)
print('chi loading...')
da = xr.open_dataarray(ifile)#.load()
chi = da

def vi2p(vi):
    """calculate TC transition probablity from ventilation index (VI) based on Tang and Emanual 2012 paper"""
    vi0 = 0.014
    n = np.log(9)/np.log(0.1/0.014)
    p = 1/(1 + (vi/vi0)**n)
    return p

#calculate VI: divide the whole period into slices and do one piece each time
print('partial VI calculating...')
time_ref = slice('1900', '1999')
if daname_partial == 'PI': #use monthly clim for Vshear and chi
    Vshear = Vshear.sel(time=time_ref).groupby('time.month').mean('time')
    chi = chi.sel(time=time_ref).groupby('time.month').mean('time')
    VI = Vshear * chi / PI.groupby('time.month')
elif daname_partial == 'Vshear': #use monthly clim for PI and chi
    PI = PI.sel(time=time_ref).groupby('time.month').mean('time')
    chi = chi.sel(time=time_ref).groupby('time.month').mean('time')
    VI = Vshear.groupby('time.month') * (chi / PI)
elif daname_partial == 'chi': #use monthly clim for PI and Vshear
    PI = PI.sel(time=time_ref).groupby('time.month').mean('time')
    Vshear = Vshear.sel(time=time_ref).groupby('time.month').mean('time')
    VI = ( Vshear / PI ) * chi.groupby('time.month')
print('partial prob calculating...')
prob = vi2p(VI) #VI to prob

#save
print('saving...')
ds = prob.drop('month').assign_attrs(long_name=f'TC transition probability from {daname_partial}').to_dataset(name='prob')
encoding = {'prob': {'zlib': True, 'complevel': 1}}
ds.to_netcdf(ofile, encoding=encoding, unlimited_dims='time')
print('[saved]:', ofile)

 
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)
 
    try:
        tt.check(f'**Done**')
    except:
        pass
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()