#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sun Sep 11 08:42:49 EDT 2022
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('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
import xfilter
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
lmr = '2019'
lpwindow = 15
func = lambda da: da.filter.lowpass(1/lpwindow, dim='year', padtype='even').sel(year=slice(850,None))

if lmr in ('2019',):
    ifile = 'LMR2019_nino34.nc'
    da = xr.open_dataarray(ifile)
    da2019 = da
else:
    ifile = 'LMR2018_nino34.nc'
    da = xr.open_dataarray(ifile)
    da2018 = da

"""
import geoxarray
#ifiles = '/tiger/scratch/gpfs/GEOCLIM/wenchang/AM2.1p1/work/amipLMR2019SST0850_e1_tigercpu_intelmpi_18_30PE/POSTP/????0101.atmos_month.nc'
ifiles = '/tiger/scratch/gpfs/GEOCLIM/wenchang/AM2.1p1/work/amipLMR2018SST0850_e1_tigercpu_intelmpi_18_30PE/POSTP/????0101.atmos_month.nc'
da = xr.open_mfdataset(ifiles)['t_surf'].load() \
    .sel(lon=slice(360-170, 360-120), lat=slice(-5,5)).geo.fldmean() \
    .groupby('time.year').mean('time') \
    .pipe(lambda da: da - da.mean('year'))
#da2019model = da
da2018model = da
"""

def wyplot(da, **kws):
    ax = kws.pop('ax', plt.gca())
    alpha = kws.pop('alpha', 0.2)
    color = kws.pop('color', None)
    #raw
    mean = da.mean('MCrun')
    mean.plot(ax=ax, color=color, **kws)
    #lowpass
    da = da.pipe(func)
    mean = da.mean('MCrun')
    mean.plot(ax=ax, color=color, label=f'{lpwindow}-year lowpass')

    ax.set_xticks(range(2000,850,-100))
    ax.set_xlim(da.year[0], da.year[-1])

def wyplot_span(*args):
    #years = 1200,1400
    ax = plt.gca()
    years = args
    ax.axvspan(*years, color='gray', alpha=0.2)

    da = da2018
    damean = da.sel(year=slice(*years)).mean(['MCrun', 'year'])
    ax.plot(years, [damean, damean], ls='--', color='C0', lw=1)
    da = da2019
    damean = da.sel(year=slice(*years)).mean(['MCrun', 'year'])
    ax.plot(years, [damean, damean], ls='--', color='C1', lw=1)

 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    figsize = 8,3.5
    fig,ax = plt.subplots(figsize=figsize)
    if lmr == '2019': label = 'LMR2.1'
    wyplot(da2019, label=label)

    #da2019model.pipe(func).plot(label='model2019')
    #da2018model.pipe(func).plot(label='model2019')

    ax.legend()

    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__LME{lmr}.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:
        pass
    else:
        plt.show()