#!/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]
    time_coords = xr.cftime_range('0000-01', '0014-12', freq='MS', calendar='noleap')#new time coordinates to be used
    tspan = slice('0005-01', '0014-12') #time-mean period
    fig, axes = plt.subplots(2,3, figsize=(10,6), sharex=True, sharey=True)
    for ax,volc in zip(axes.flat, da.volc.values):
        #annual mean clim from control
        levels = np.arange(12,31,4)
        cs = da_ref.mean('month').plot.contour(ax=ax, yincrease=False, colors='k', alpha=1/4, levels=levels)
        ax.clabel(cs, cs.levels[0::2])
        #control + change
        #cs = da_ref.mean('month').pipe(lambda x: x + damean).plot.contour(ax=ax, yincrease=False, colors='k', alpha=1/4, levels=levels, linestyles='dashed')
        #ax.clabel(cs, cs.levels[0::2])
        #mean
        damean = da.sel(volc=volc).assign_coords(time=time_coords).sel(time=tspan).mean('time').mean('en')
        im = damean.plot.contourf(ax=ax, yincrease=False, levels=np.arange(-0.6,0.61,0.1), extend='both', add_colorbar=False)
        #significance
        ds = da.sel(volc=volc).assign_coords(time=time_coords).sel(time=tspan).mean('time').pipe(cim, dim='en')
        damean.where(np.abs(damean)>ds.err).pipe(lambda x: x*0).plot.contourf(ax=ax, colors='none', hatches=['....'], add_colorbar=False, yincrease=False)


    cbar_label = f'{tspan.start} to {tspan.stop} temp response [K]'
    fig.colorbar(im, ax=axes, label=cbar_label)

    for ax in axes.flat:
        ax.set_xlabel('')
    for ax in axes[:, 1:].flat:
        ax.set_ylabel('')
    for ax in axes[-1, :]:
        xticks2lat(np.arange(-30,31,10), ax=ax, rotation=30)

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