#!/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

ifile = 'LMRseasonal_nino34.nc'
da = xr.open_dataarray(ifile)
da_seasonal = da.groupby('time.year').mean('time')

"""
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)
    if 'MCrun' in da.dims:
        mean = da.mean('MCrun')
        std = da.std('MCrun')
        ax.fill_between(da.year, mean-std, mean+std, color=color, alpha=alpha)
    else:
        mean = da
    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)
    da = da_seasonal
    damean = da.sel(year=slice(*years)).mean('year')
    ax.plot(years, [damean, damean], ls='--', color='C2', 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='LMR2.0')#2018')
    wyplot(da2019, color='C1', label='LMR2.1')#2019')
    wyplot(da_seasonal, label='LMR_seasonal')
    #da_seasonal.pipe(func).plot(label='LMR_seasonal')

    #wyplot_span(1200, 1400)
    #wyplot_span(1450, 1550)
    #wyplot_span(1551, 1650)
    wyplot_span(1351, 1550)
    wyplot_span(1701, 1900)

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

    ax.legend()
    ax.set_title(f'{lpwindow}-year lowpass Nino3.4 anomalies')

    
    #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()