#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed May 13 06:26:13 PM EDT 2026
import sys
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
if __name__ == '__main__':
    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
from misc import get_kws_from_argv
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
#settings
ens = get_kws_from_argv('ens', 3)
ens = int(ens)
year = f'{ens:04d}'
selregion = lambda x: x.sel(lat=slice(10,80), lon=slice(200,360))
ifile = '/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/CTL9120/4xdaily2/0001/autoregressive_predictions.nc'
#process
ds = xr.open_dataset(ifile)
ds = ds.squeeze()
time= xr.cftime_range('0001-01-01', freq='6H', periods=ds.time.size, calendar='noleap')
ds['time'] = time
ts = ds.PRATEsfc.sel(lat=40.3504, lon=360-74.6571, method='nearest').load()
ts = (ts*24*3600).assign_attrs(units='mm/day')
imax = ts.sel(time=year).argmax('time').item()
h500 = ds.h500.sel(time=year).isel(time=imax)

pr = ds.PRATEsfc.sel(time=year).isel(time=imax)
pr = (pr*24*3600).assign_attrs(units='mm/day')

#plot
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    xr.set_options(keep_attrs=True)
    if 'line' in sys.argv:
        fig,ax = plt.subplots()
        ts.groupby('time.year').max('time').plot()
        #savefig
        if 'savefig' in sys.argv or 's' in sys.argv:
            figname = __file__.replace('.py', f'__line.png')
            if 'overwritefig' in sys.argv or 'o' in sys.argv:
                wysavefig(figname, overwritefig=True)
            else:
                wysavefig(figname)

    fig,ax = plt.subplots()
    pr.pipe(selregion).plot(cmap='YlGnBu', ax=ax) 
    cs = (h500 - h500.mean('lon')).pipe(selregion) \
        .plot.contour(colors='C3', levels=20, linewidths=0.5, ax=ax)
    ax.clabel(cs, cs.levels[::2])
    mapplot()
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__{year}.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()