#!/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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
thisdir = os.path.dirname(__file__)
#global mean Ts
"""
ibasename = 't_surf.glbmean.volcs.nens30.nc'
ifile = os.path.join(thisdir, 'data', ibasename)
da = xr.open_dataarray(ifile).load()
da = da.groupby('time.month') - da.sel(volc='CTL').mean('en').groupby('time.month').mean('time')
da_glbmean = da
"""
#NH-SH Ts
ibasename = 't_surf.dHmean.volcs.nens30.nc'
ifile = os.path.join(thisdir, 'data', ibasename)
da = xr.open_dataarray(ifile).load()
da = da.groupby('time.month') - da.sel(volc='CTL').mean('en').groupby('time.month').mean('time')
#da_dHmean = da
#da = da_dHmean/2/da_glbmean.mean('en').mean('time') #asymmetry: (NH-SH)/(NH+SH)
#remove ensemble-mean climatology from the control
#time_new = xr.cftime_range('1001-01-01', periods=12*15, freq='MS')#, calendar='noleap')
#da = da.assign_coords(time=time_new)
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, ax = plt.subplots()
    n_window = 5 # 5 months in the filter window 
    for volc in da.volc.values:
        if volc == 'CTL':
            color = 'k'
        else:
            color = None
        da.sel(volc=volc).mean('en') \
            .filter.lowpass(1/n_window, dim='time', padtype='even') \
            .plot(label=volc, color=color)

    #ax.set_ylim(-1,1)
    ax.set_ylabel('NH$-$SH t_surf [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-2000:02d}' for yr in times.dt.year.values])#da.time.dt.year.values[0::24]])
    ax.set_xlabel('year')
    ax.legend(ncol=3)
    

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