#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Jun 18 10:44:54 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 misc.cim import cim
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'data/t_surf.CTL1860_noleap_tigercpu_intelmpi_18_576PE.atmos_month.nens30.nyears15.nc'
da_ctl = xr.open_dataarray(ifile).assign_coords(time=xr.cftime_range('0001-01', '0015-12', freq='MS'))
da_ctl.load()
volc = 'StMaria'
ifile = f'data/t_surf.{volc}_PI_ens_noleap.atmos_month.nens30.nc'
da_volc = xr.open_dataarray(ifile).assign_coords(time=xr.cftime_range('0001-01', '0015-12', freq='MS'))
da_volc.load()
#difference between two periods
#func_diff = lambda x: x.sel(time=slice('0010-12', '0011-02')).mean('time') - x.sel(time=slice('0008-12', '0009-02')).mean('time')
tspan = slice('0013-01', '0013-12')
func_diff = lambda x: x.sel(time=tspan).mean('time') - da_ctl.sel(time=tspan).mean('time').mean('en')
da = da_volc.pipe(func_diff)
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig, ax = plt.subplots()
    damean = da.mean('en')
    da.mean('en').plot.contourf(levels=21, ax=ax, robust=True)
    err = da.pipe(cim, dim='en')
    damean.where(np.abs(damean)>err).pipe(lambda x: x*0).plot.contourf(colors='none', add_colorbar=False, hatches=['////'])
    #plt.sca(ax)
    mapplot(ax=ax, fill_continents=True)
    ax.set_title(f'{volc} Tsa over {tspan.start} to {tspan.stop}') 

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