#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue May 28 11:42:33 EDT 2024
if __name__ == '__main__':
    import sys,os
    from misc.timer import Timer
    tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
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
from misc.seasons import sel_season
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
def wyplot_basin_boundary(ax=None, **kws):
    latmax = 60
    if ax is None: fig,ax = plt.subplots()
    color = kws.pop('color', 'C2') #default color set to 'C2'
    textcolor = 'gray'
    #NA
    xy = np.array([[360-100, latmax], [360-100, 20],[360-65, 0], [360, 0], [360, latmax] ])
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    ax.text(360-100, 20, '100W', ha='right', va='bottom', color=textcolor)
    ax.text(360-65, 0, '65W', ha='left', va='bottom', color=textcolor)
    #EA
    xy = np.array([[360-160, latmax], [360-160, 0], [360-65, 0]] )
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    ax.text(360-160, 0, '160W', ha='left', va='bottom', color=textcolor)
    #WP
    xy = np.array([[105, latmax], [105, 0], [360-160, 0]] )
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    ax.text(105, 0, '105E', ha='left', va='bottom', color=textcolor)
    #NI
    xy = np.array([[30, latmax], [30, 0], [105, 0]] )
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    #SI
    xy = np.array([[30, -latmax], [30, 0], [105, 0], [105, -latmax]] )
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    #AU
    xy = np.array([[105, 0], [165, 0], [165,-latmax]] )
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    ax.text(165, -1, '165E', ha='left', va='top', color=textcolor)
    #SP
    xy = np.array([[165, 0], [360-70, 0], [360-70,-latmax]] )
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    ax.text(360-70, -1, '70W', ha='right', va='top', color=textcolor)
    ##SA
    #xy = np.array([[360-70, 0], [360, 0], [360,-latmax]] )
    #ax.plot(xy[:,0], xy[:,1], color=color, **kws)

#nino and nina are based on year0 from https://origin.cpc.ncep.noaa.gov/products/analysis_monitoring/ensostuff/ONI_v5.php
nino_years = [1982, 1986, 1987, 1991, 1994, 1997, 2002, 2009, 2015]
nina_years = [1988, 1995, 1998, 1999, 2007, 2010, 2011, 2017]
nh_season = 'MJJASON'
sh_season = 'ONDJFMAM'
ifile = 'IBTrACS.ALL.v04r00.2022-01-25.density.TS17.1980-2021.nc'
da = xr.open_dataset(ifile)['density'].load()
units = da.attrs['units']
def do_anom(years):
    #NH season
    daa_nh = da.pipe(sel_season, nh_season).resample(time='AS-MAY').mean('time') \
        .groupby('time.year').mean('time') \
        .pipe(lambda x: x.sel(year=years).mean('year') - x.mean('year'))
    #SH season
    daa_sh = da.pipe(sel_season, sh_season).resample(time='AS-OCT').mean('time').isel(time=slice(1,-1)) \
        .groupby('time.year').mean('time') \
        .pipe(lambda x: x.sel(year=years).mean('year') - x.mean('year'))
    #merge NH and SH
    daa = daa_nh.where(daa_nh.lat>=0, other=daa_sh).assign_attrs(units=units)

    return daa
daa_nino = do_anom(nino_years)
daa_nina = do_anom(nina_years)


 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig,axes = plt.subplots(2,1, figsize=(8,6))
    levels = np.arange(-1.3, 1.31, 0.2)
    cmap = 'seismic'
    cmap = 'RdBu_r'
    func = lambda x: x.sel(lat=slice(-60,60))
    ax = axes[0]
    im = daa_nino.pipe(func).plot.contourf(levels=levels, ax=ax, add_colorbar=False, cmap=cmap)
    wyplot_basin_boundary(ax=ax)
    ax = axes[1]
    im = daa_nina.pipe(func).plot.contourf(levels=levels, ax=ax, add_colorbar=False, cmap=cmap)
    wyplot_basin_boundary(ax=ax)

    for ax in axes:
        plt.sca(ax)
        ax.set_xlim(0,360)
        ax.set_xticks(np.arange(0,361,60))
        mapplot()
        ax.set_xlabel('')
        ax.set_ylabel('')
        #ax.set_title('')
        #ax.set_ylim(-60,60)
    axes[0].set_title('(a) IBTrACS El Niño - clim')
    axes[1].set_title('(b) IBTrACS La Niña - clim')
    fig.colorbar(im, ax=axes, label=units, ticks=levels, aspect=30) 
    

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