#!/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
lpwindow = 15
func = lambda da: da.filter.lowpass(1/lpwindow, dim='year', padtype='even').sel(year=slice(850,None))

ifile = 'LMR2019_nino34.nc'
da = xr.open_dataarray(ifile)
da2019 = da

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)
    da = da.pipe(func)
    mean = da.mean('MCrun')
    std = da.std('MCrun')
    ax.fill_between(da.year, mean-std, mean+std, color=color, alpha=alpha)
    mean.plot(ax=ax, color=color, **kws)

    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)
    wyplot(da2018, color='C0', label='LMR2018')
    wyplot(da2019, color='C1', label='LMR2019')

    wyplot_span(1200, 1400)
    wyplot_span(1450, 1550)
    wyplot_span(1551, 1650)

    #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'__{lpwindow}lowpass.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()