#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Feb 22 13:49:21 EST 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
from misc.cim import cim
import xfilter
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
daname = 't_surf'
ifile = f'{daname}.glbmean.volcs.nens10.nc'
da = xr.open_dataarray(ifile).load()
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #import nc_time_axis
    fig, ax = plt.subplots()
    alpha = 0.3
    volcs = da.volc.values[1:]
    func_anom = lambda da: da.sel(volc=volcs) - da.sel(volc='CTL').mean('en')
    daa = da.pipe(func_anom)
    #ignore values before eruption date
    #emonth = {'PInatubo': 6, 'Agung': 3, 'StMaria': 10, 'Chichon': 3, 'Chichon_CMIP5volc': 3}
    emonths = xr.DataArray([6, 3, 10, 3, 3], dims='volc', coords=[volcs])
    months = xr.DataArray(np.vstack([np.arange(1, daa.time.size+1),]*daa.volc.size), dims=['volc', 'time'], coords=[volcs, daa.time])
    daa = daa.where(months>emonths, other=0)
    #print(daa.isel(time=slice(0,12), en=0))
    #lowpass filter
    #daa = daa.filter.lowpass(1/9, dim='time', padtype='odd')
    daa_cim = cim(daa, dim='en')
    daa_mean = daa.mean('en')
    for volc in volcs:
        y1 = daa_mean.sel(volc=volc) - daa_cim.sel(volc=volc)
        y2 = daa_mean.sel(volc=volc) + daa_cim.sel(volc=volc)
        plt.fill_between(daa_mean.time.values, y1, y2, alpha=alpha)
        daa_mean.sel(volc=volc).plot(label=volc, lw=2)
    ax.axhline(0, color='gray', ls='--')
    ax.legend()
    ax.set_title('')

    
    figname = __file__.replace('.py', f'_{tt.today()}.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        plt.savefig(figname)
        print('[saved]:', figname)
    tt.check(f'**Done**')
    plt.show()