#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Feb 12 04:16:24 PM EST 2025
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
#eta: \eta = f + \zeta
#ifile = 'eta_AM2.5C360_amipHadISSTrcp45_tigercpu_intelmpi_18_1080PE_ens16_1871-2100.nc'
ifile = glob.glob('eta_*2021.nc')[0]
print('ifile:', ifile)
#ofile = 'spi_AM2.5C360_amipHadISSTrcp45_tigercpu_intelmpi_18_1080PE_ens16_1871-2100.nc'
ofile = ifile.replace('eta_', 'spi_')
if os.path.exists(ofile): print('[exists]:', ofile); sys.exit()
da = xr.open_dataarray(ifile).load()#.rename(grid_xt='lon', grid_yt='lat')
eta = da
#omg500
#ifile = 'omega500_AM2.5C360_amipHadISSTrcp45_tigercpu_intelmpi_18_1080PE_ens16_1871-2100.nc'
ifile = glob.glob('omega500_*2021.nc')[0]
print('ifile:', ifile)
da = xr.open_dataset(ifile)['omega'].isel(level=0).load()#.rename(grid_xt='lon', grid_yt='lat')
omg500 = da

print('ofile:', ofile)
print('spi...')
#eta_y
alpha = 0.69 # used in 1/(1 + Z^{-1/\alpha})
R = 6370*1000# Earth's radius in unites of m
U = 20 #m/s, constant
y = R * eta.lat/180*np.pi #y coordinate to replace "lat"
#use y coordinate (units of m) to do the differentiation and convert back to "lat" after that
eta_y = eta.rename(lat='y').assign_coords(y=y.values).differentiate('y').rename(y='lat').assign_coords(lat=eta.lat.values)
#Z value in Tsung-Lin's paper: eta uses absolute value of cyclonic grids and zero for anticyclonic grids
Z = eta.where(eta.lat>0, other=-eta).clip(min=0) / np.sqrt( np.abs(eta_y)*U )
#spi
spi = (-omg500).clip(min=0)/(1 + Z**(-1/alpha))

#save
print('save...')
ds = spi.to_dataset(name='spi')
ds.to_netcdf(ofile, encoding={'spi': {'zlib': True, 'complevel': 1}})
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()