#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Jun 21 21:14: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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'data/temp.x170W120Wmean.volcs.nens30.40S-40N_0-300m.nc'
da = xr.open_dataarray(ifile).load()
da_ref = da.sel(volc='CTL').mean('en').groupby('time.month').mean('time')
da = da.groupby('time.month') - da_ref

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import xticks2lat
    from xcim import cim
    volc = ['Agung','Chichon', 'Chichon_CMIP5volc','Pinatubo','StMaria'][1]
    """
    da_ = da.sel(volc=volc).mean('en').groupby('time.year').mean('time').assign_coords(year=range(15))
    g = da_.plot.contourf(col='year', col_wrap=1, yincrease=False, robust=False, levels=np.arange(-1,1.01,0.1), figsize=(8,6))
    for ax,year in zip(g.axes.flat, da_.year.values):
        ax.set_title('')
        ax.text(0,1,f'{year = }', transform=ax.transAxes, ha='left', va='top')
        ax.set_ylabel('m')
        ax.set_ylim(200,0)
        ax.set_xlabel('')
    xticks2lat(ax=g.axes.flat[-1])
    """
    ds = da.sel(volc=volc).groupby('time.year').mean('time').assign_coords(year=range(15)).pipe(cim, dim='en')
    fig, axes = plt.subplots(5,3, figsize=(10,6), sharex=True, sharey=True)
    for ax,year in zip(axes.transpose().flat, ds.year.values):
        im = ds.mean_value.sel(year=year).plot.contourf(yincrease=False, levels=np.arange(-1,1.01,0.1), ax=ax, add_colorbar=False)
        ds.mean_value.where(np.abs(ds.mean_value)>ds.err).pipe(lambda x: x*0).sel(year=year) \
            .plot.contourf(ax=ax, colors='none', hatches=['....'], add_colorbar=False, yincrease=False)
        cs = da_ref.mean('month').plot.contour(ax=ax, yincrease=False, colors='k', alpha=1/4)
        ax.clabel(cs, cs.levels[0::2])

        ax.set_title('')
        ax.text(0,1,f'{ year = }', transform=ax.transAxes, ha='left', va='top')
        ax.set_xlabel('')
        ax.set_ylabel('')
    for ax in axes.flat[-3:]:
        xticks2lat(np.arange(-30, 31, 10), ax=ax, rotation=30)
    for ax in axes[:,0]:
        ax.set_ylabel('depth [m]')
        #ax.set_ylim(200,0)
    fig.colorbar(im, ax=axes, label=f'{volc} temp response [K]')

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