#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Apr 15 09:56:26 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
import xfilter
from misc.cim import cim
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
thisdir = os.path.dirname(__file__)
ibasename = 'nino34.volcs.nens30.nc'
ifile = os.path.join(thisdir, 'data', ibasename)
da = xr.open_dataarray(ifile).load()
#time_new = xr.cftime_range('1001-01-01', periods=12*15, freq='MS')#, calendar='noleap')
#da = da.assign_coords(time=time_new)
 
def wyplot(ax=None, volc='Pinatubo', color='C0'):
    if ax is None:
        fig, ax = plt.subplots()
    alpha = 0.2
    n_window = 5 # 5 months in the filter window 
    smooth = lambda x: x.filter.lowpass(1/n_window, dim='time', padtype='even')
    #volc
    #volc = 'Pinatubo'
    #color = 'C0'
    damean = da.sel(volc=volc).pipe(smooth).mean('en')
    spread = da.sel(volc=volc).pipe(smooth).pipe(cim, dim='en')
    ax.fill_between(damean.time, damean-spread, damean+spread, alpha=alpha, color=color)
    damean.plot(ax=ax, label=volc, color=color)

    ax.set_ylim(-1.5,1.5)
    ax.set_yticks(np.arange(-1.5,1.51,0.5))
    ax.set_ylabel('Nino3.4 [K]')
    #ax.spines['right'].set_visible(False)
    #ax.spines['top'].set_visible(False)
    times = da.time[0::12]
    ax.set_xlim(da.time[0], da.time[-1])
    ax.set_xticks(times)
    ax.set_xticklabels([f'{yr-2001:02d}' for yr in times.dt.year.values])#da.time.dt.year.values[0::24]])
    ax.set_xlabel('year')
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, axes = plt.subplots(3, 2, sharex=True, sharey=True, figsize=(8,6))
    #volcs = da.volc.values
    volcs = np.array(['CTL', 'StMaria', 'Pinatubo', 'Chichon', 'Agung', 'Chichon_CMIP5volc'])
    colors = ['gray'] + [f'C{ii}' for ii in range(volcs.size)]
    for ax,volc,color in zip(axes.flat, volcs, colors):
        wyplot(ax=ax, volc=volc, color=color)
    
    for ax in axes.flat:
        ax.set_title('')
        ax.legend(loc='upper left')
    for ax in axes[:,1]:
        ax.set_ylabel('')
    for ax in axes.flat[:-2]:
        ax.set_xlabel('')

    #savefig
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()