#!/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 yticks2lat
    from xcim import cim
    #volc = ['Agung','Chichon', 'Chichon_CMIP5volc','Pinatubo','StMaria'][1]
    time_coords = pd.date_range('2000-01', '2014-12', freq='MS')#new time coordinates to be used
    zspan = slice(50,200)
    fig, axes = plt.subplots(3,2, figsize=(10,6), sharex=True, sharey=True)
    for ax,volc in zip(axes.flat, da.volc.values):
        damean = da.sel(volc=volc).assign_coords(time=time_coords).sel(st_ocean=zspan).mean('st_ocean').mean('en')
        levels = np.arange(-1.2,1.21,0.2)
        add_colorbar = True if levels is None else False
        im = damean.plot.contourf(ax=ax, levels=levels, extend='both', add_colorbar=add_colorbar, y='yt_ocean')
        #significance
        ds = da.sel(volc=volc).assign_coords(time=time_coords).sel(st_ocean=zspan).mean('st_ocean').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, y='yt_ocean')


    cbar_label = f'{zspan.start}-{zspan.stop}m mean temp response [K]'
    fig.colorbar(im, ax=axes, label=cbar_label)

    ax.set_ylim(-30,30)
    for ax in axes[:-1, :].flat:
        ax.set_xlabel('')
    for ax in axes.flat:
        ax.set_ylabel('')
    for ax in axes[:, 0]:
        yticks2lat(np.arange(-30,31,10), ax=ax)#, rotation=30)
    for ax in axes[-1, :]:
        xticks = [f'{yyyy:04d}-01' for yyyy in range(2000,2015)]
        xticklabels = [f'{yy:02d}' for yy in range(15)]
        ax.set_xticks(xticks)
        ax.set_xticklabels(xticklabels)

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