#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Jun 18 16:53:32 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
from xcim import cim
from geoplots import yticks2lat
import xfilter
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'data/t_surf.xmean.volcs.nens30.nc'
da = xr.open_dataarray(ifile)
da = da.groupby('time.month') - da.sel(volc='CTL').mean('en').groupby('time.month').mean('time')
da = da.drop('time')#.assign_coords(time=xr.cftime_range('0001-01', '0015-12', freq='MS'))
 
 
def wyplot(ax=None, volc='StMaria', **kws):
    if ax is None:
        fig, ax = plt.subplots()
        yticks = range(-80,81,20)
    else:
        yticks = None
        
    smooth = lambda x: x.filter.lowpass(1/5, dim='time', padtype='even')
    #volc = 'StMaria'
    ds = da.sel(volc=volc).pipe(smooth).pipe(cim, dim='en')
    im = ds.mean_value.plot.contourf(y='lat', ax=ax, robust=True, levels=21, vmax=1, extend='both', **kws)
    ds.mean_value.where(np.abs(ds.mean_value)>ds.err).pipe(lambda x: x*0) \
        .plot.contourf(y='lat', ax=ax, colors='none', hatches=['////'], add_colorbar=False)

    ax.set_xticks(range(0,12*15,12))
    ax.set_xticklabels(range(15))
    if yticks is not None:
        yticks2lat(range(-80,81,20), ax=ax)
    ax.set_ylabel('lat')
    ax.set_xlabel('year')

    return im

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    volcs = ['CTL', 'StMaria',
        'Pinatubo', 'Chichon',
        'Agung', 'Chichon_CMIP5volc'
        ]
    fig, axes = plt.subplots(3, 2, figsize=(10,6), sharex=True, sharey=True)
    for ax,volc in zip(axes.flat, volcs):
        im = wyplot(ax=ax, volc=volc, add_colorbar=False)

    ax = axes[0,0]
    ax.set_yticks(range(-90,91,30))
    yticks2lat(ax=ax)
    fig.colorbar(im, ax=axes, label='Ts response [K]', shrink=0.8)
    for ax in axes.flat[:4]:
        ax.set_xlabel('')
    for ax in axes[:,1]:
        ax.set_ylabel('')

    """
    #each volc in a single figure
    for volc in volcs:
        wyplot(volc=volc) 
    """

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