#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Jun 23 18:08:42 EDT 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'data/nino34.volcs.nens1.nyears100.nc'
da = xr.open_dataarray(ifile).sel(en=1)#.sel(time=slice('2005', None))
 
def wyplot(ax=None, nyears_ignore=0):
    if ax is None:
        fig, ax = plt.subplots()
    da_ = da.sel(time=slice(f'{2001+nyears_ignore}', None))
    nyears_rolling = 30
    da_ = da_.rolling(time=nyears_rolling*12+1, center=True).std()
    colors = ['k',] + [f'C{ii}' for ii in range(5)]
    for volc,color in zip(da.volc.values, colors):
        da_.sel(volc=volc).plot(color=color, label=volc, ax=ax)

    ax.legend(ncol=3, loc='lower center')
    ax.set_title(f'{nyears_rolling}-year running STD of Nino3.4 index')
    ax.set_ylabel('K')
    xticks = [f'{yyyy:04d}-01' for yyyy in range(2001,2001+100, 10)]
    xticklabels = [f'{yy:02d}' for yy in range(0,100,10)]
    ax.set_xticks(xticks)
    ax.set_xticklabels(xticklabels)

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import xfilter
    #smooth = lambda x: x.filter.lowpass(1/12/5, dim='time', padtype='even')
    #da.pipe(smooth).plot(hue='volc')
    fig, axes = plt.subplots(1, 2, sharey=True, figsize=(10,4))
    for ax,n in zip(axes, [0, 4]):
        wyplot(ax=ax, nyears_ignore=n)
    ax = axes[-1]
    ax.text(1, 1, f'first {n} years ignored', transform=ax.transAxes, ha='right', va='top')

    
    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()