#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Feb 20 11:01:22 AM EST 2026
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
case = get_kws_from_argv('case', 'base')
nyears = get_kws_from_argv('nyears', 1); nyears = int(nyears)
ofile = __file__.replace('.py', f'__{case}_{nyears}years.nc')
if os.path.exists(ofile):
    stp = xr.open_dataarray(ofile)
    print('[loaded]:', ofile)
else:
    ifiles = {
    'cape': 'cape_AM2.5C360_CTL2010s_tiger3_intel24ifort_openmpi_2430PE_0101-0150_atmos_daily.nc',
    'cin':  'cin_AM2.5C360_CTL2010s_tiger3_intel24ifort_openmpi_2430PE_0101-0150_atmos_daily.nc',
    'srh01':  'srh01_AM2.5C360_CTL2010s_tiger3_intel24ifort_openmpi_2430PE_0101-0150_atmos_daily.nc',
    'lcl':  'data_AM2.5C360_CTL2010s_tiger3_intel24ifort_openmpi_2430PE_0101-0150_atmos_daily_sblcl.nc',
    'bwd06':  'data_AM2.5C360_CTL2010s_tiger3_intel24ifort_openmpi_2430PE_0101-0150_atmos_daily_bwd.nc',
    }
    ifile = ifiles['cape']
    da = xr.open_dataarray(ifile).sel(time=slice(f'{150-nyears+1:04d}', None))
    #print(da); sys.exit()
    sbcape = da
    ifile = ifiles['cin']
    da = xr.open_dataarray(ifile).sel(time=slice(f'{150-nyears+1:04d}', None))
    #print(da); sys.exit()
    sbcin = da
    ifile = ifiles['lcl']
    da = xr.open_dataarray(ifile).sel(time=slice(f'{150-nyears+1:04d}', None))
    #print(da); sys.exit()
    sblcl = da
    ifile = ifiles['srh01']
    da = xr.open_dataarray(ifile).sel(time=slice(f'{150-nyears+1:04d}', None))
    #print(da); sys.exit()
    srh01 = da.where(da.lat>=0, other=-da)
    ifile = ifiles['bwd06']
    da = xr.open_dataarray(ifile).sel(time=slice(f'{150-nyears+1:04d}', None))
    #print(da); sys.exit()
    bwd06 = da
    
    if case == 'base':
        stp = (sbcape/1500).clip(0) * ((2000 - sblcl)/1000).clip(0) * (srh01/150).clip(0) * (bwd06/20).clip(0) * ((200-sbcin)/150).clip(0, 1)
    elif case == 'noCIN':
        stp = (sbcape/1500).clip(0) * ((2000 - sblcl)/1000).clip(0) * (srh01/150).clip(0) * (bwd06/20).clip(0)
    elif case == 'switchCIN':
        stp = (sbcape/1500).clip(0) * ((2000 - sblcl)/1000).clip(0) * (srh01/150).clip(0) * (bwd06/20).clip(0)
        stp = stp.where(sbcin<50, other=0)
    elif case == 'switchCIN100':
        stp = (sbcape/1500).clip(0) * ((2000 - sblcl)/1000).clip(0) * (srh01/150).clip(0) * (bwd06/20).clip(0)
        stp = stp.where(sbcin<100, other=0)
    elif case == 'switchCIN150':
        stp = (sbcape/1500).clip(0) * ((2000 - sblcl)/1000).clip(0) * (srh01/150).clip(0) * (bwd06/20).clip(0)
        stp = stp.where(sbcin<150, other=0)
    elif case == 'switchCIN200':
        stp = (sbcape/1500).clip(0) * ((2000 - sblcl)/1000).clip(0) * (srh01/150).clip(0) * (bwd06/20).clip(0)
        stp = stp.where(sbcin<200, other=0)
    elif case == 'switchCIN300':
        stp = (sbcape/1500).clip(0) * ((2000 - sblcl)/1000).clip(0) * (srh01/150).clip(0) * (bwd06/20).clip(0)
        stp = stp.where(sbcin<300, other=0)
    elif case == 'cin100':
        stp = (sbcape/1500).clip(0) * ((2000 - sblcl)/1000).clip(0) * (srh01/150).clip(0) * (bwd06/20).clip(0) * ((250-sbcin)/150).clip(0, 1)
    elif case == 'cin150':
        stp = (sbcape/1500).clip(0) * ((2000 - sblcl)/1000).clip(0) * (srh01/150).clip(0) * (bwd06/20).clip(0) * ((300-sbcin)/150).clip(0, 1)
    elif case == 'cin200':
        stp = (sbcape/1500).clip(0) * ((2000 - sblcl)/1000).clip(0) * (srh01/150).clip(0) * (bwd06/20).clip(0) * ((350-sbcin)/150).clip(0, 1)
    elif case == 'cin300':
        stp = (sbcape/1500).clip(0) * ((2000 - sblcl)/1000).clip(0) * (srh01/150).clip(0) * (bwd06/20).clip(0) * ((450-sbcin)/150).clip(0, 1)
    stp = stp.clip(0,1)

    #save
    ds = stp.to_dataset(name='stp')
    encoding = {'stp': {'dtype': 'float32', 'zlib': True, 'complevel': 1}}
    ds.to_netcdf(ofile, encoding=encoding)
    print('[saved]:', ofile)
stp = stp.mean('time')
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    from misc.landmask import whereland
    da = stp
    da.pipe(whereland).plot(robust=True, vmax=0.01)
    mapplot()

    ax = plt.gca()
    title = 'stp'
    if case != 'base': title += ' ' + case
    ax.set_title(title)

    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__{case}_{nyears}years.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()