#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Oct 16 11:07:13 AM EDT 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
import geoxarray, xfilter
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
case = get_kws_from_argv('case', 'all')

#HadISST
ifile = '/projects/w/wenchang/analysis/TC_statmodel/HadISST.sstindex.187001-201912.nc'
ds = xr.open_dataset(ifile).load().groupby('time.year').mean('time').sel(year=slice(None, 2000))
mdr_obs = ds.sst_mdr
mdr_obs -= mdr_obs.sel(year=slice(1950,1980)).mean('year')
trop_obs = ds.sst_trop
trop_obs -= trop_obs.sel(year=slice(1950,1980)).mean('year')
rmdr_obs = mdr_obs - trop_obs

#LMR seasonal
#ifile = '../tos_mean.nc'
ifile = '/projects/w/wenchang/data/LMR/seasonal/tos_mean.nc'
da = xr.open_dataarray(ifile).load()
da = da.groupby('time.year').mean('time')
#trop
trop = da.sel(lat=slice(-30,30)).geo.fldmean()
#mdr
mdr = da.sel(lat=slice(10,25), lon=slice(280,340)).geo.fldmean()
rmdr = mdr - trop

#LMR2.1
ifile = '/tigress/gvecchi/DATA/LMR_2019/sst_MCruns_ensemble_mean.nc'
da = xr.open_dataarray(ifile).load()
da = da.groupby('time.year').mean('time').mean('MCrun')
#trop
trop_lmr2p1 = da.sel(lat=slice(-30,30)).geo.fldmean()
#mdr
mdr_lmr2p1 = da.sel(lat=slice(10,25), lon=slice(280,340)).geo.fldmean()
rmdr_lmr2p1 = mdr_lmr2p1 - trop_lmr2p1
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    lpwindow = 40 
    func = lambda da: da.filter.lowpass(1/lpwindow, dim='year', padtype='even')#.sel(year=slice(850,None))

    if case in ('trop', 'all'):
        trop_obs.pipe(func).plot(label='TROP mean SST from HadISST', color='gray', lw=3)
    if case in ('mdr', 'all'):
        mdr_obs.pipe(func).plot(label='MDR mean SST from HadISST', color='k', lw=3)
    if case in ('rmdr', 'all'):
        rmdr_obs.pipe(func).plot(label='rMDR mean SST from HadISST', color='k',lw=3, ls='--')
    if case in ('trop', '2.1', 'all'):
        trop_lmr2p1.plot(color='C0', lw=0.5, alpha=0.5)
        trop_lmr2p1.pipe(func).plot(label='TROP mean SST from LMR2.1', color='C0')
    if case in ('mdr', '2.1', 'all'):
        mdr_lmr2p1.plot(color='C1', lw=0.5, alpha=0.5)
        mdr_lmr2p1.pipe(func).plot(label='MDR mean SST from LMR2.1', color='C1')
    if case in ('rmdr', '2.1', 'all'):
        rmdr_lmr2p1.plot(color='C2', lw=0.5, alpha=0.5)
        rmdr_lmr2p1.pipe(func).plot(label='rMDR mean SST from LMR2.1', color='C2')

    if case in ('trop', 'seasonal', 'all'):
        trop.plot(color='C3', lw=0.5, alpha=0.5)
        trop.pipe(func).plot(label='TROP mean SST from LMR_seasonal ', color='C3')
    if case in ('mdr', 'seasonal', 'all'):
        mdr.plot(color='C4', lw=0.5, alpha=0.5)
        mdr.pipe(func).plot(label='MDR mean SST from LMR_seasonal', color='C4')
    if case in ('rmdr', 'seasonal', 'all'):
        rmdr.plot(color='C5', lw=0.5, alpha=0.5)
        rmdr.pipe(func).plot(label='rMDR mean SST from LMR_seasonal', color='C5')


    ax = plt.gca()
    ax.legend()
    ax.set_xlim(800, None)
    ax.set_ylabel('K')
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__{case}.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()