#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Oct  5 11:09:59 EDT 2022
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('start ' + ' '.join(sys.argv))
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
from modelout import get_modelout_data, update_modelout_data
from modelout.getdata import funcs
import xfilter
nwindow, dimlp = 9, 'year'
#lowpass = lambda x: x.filter.lowpass(1/nwindow, dim=dimlp, padtype='odd')
lowpass = lambda x: x.rolling(year=nwindow, center=True, min_periods=1).mean() if x.year.size>9 else x
import geoxarray
from misc import get_kws_from_argv
from misc.seasons import sel_season
from misc import roll_lon
from misc.comp2samp import comp2samp
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
daname = get_kws_from_argv('daname', 'slp')
#func = lambda x: x.load().geo.fldmean()
#funcname = 'glbmean'
#func = funcs['nino34']
#funcname = 'nino34'
season = 'DJF'
funcname = f'{season}mean'
func = lambda da: da.pipe(sel_season, season).load().resample(time='AS').mean(keep_attrs=True) #seasonmean

cleanup = False #remove cache or not

das = []
labels = []

model = 'AM2.5'
#ctl
label = f'{model}_ctl'
expname = 'CTL1990s_tigercpu_intelmpi_18_540PE'
da = update_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)
labels.append(label)
das.append(da)
#plusColdBlob
label = f'{model}_plusColdBlob'
expname = 'CTL1990s_plusColdBlob_tigercpu_intelmpi_18_540PE'
da = update_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)#, years=range(100,201))
labels.append(label)
das.append(da)

model = 'HIRAM'
#ctl
label = f'{model}_ctl'
expname = 'CTL1990s_v201910_tigercpu_intelmpi_18_540PE'
da = update_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)#, years=range(100,201))
labels.append(label)
das.append(da)
#plusColdBlob
label = f'{model}_plusColdBlob'
expname = 'CTL1990s_v201910_plusColdBlob_tigercpu_intelmpi_18_540PE'
da = update_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)#, years=range(100,201))
labels.append(label)
das.append(da)

#units
units = da.attrs['units']
 
def wyplot(da, label, ax=None, **kws):
    if ax is None: ax = plt.subplots()[1]
    #da = da.groupby('time.year').mean('time')
    da_ref = das[0] if 'AM2.5' in label else das[2]
    daa = da.mean('time') - da_ref.mean('time') #anomal from the original control
    #da = da.pipe(lowpass)
    #da = da.sel(year=slice(1900,2000))
    daa = daa.pipe(roll_lon)
    daa.assign_attrs(units=units).plot.contourf(ax=ax, **kws)

    #significance info
    ds = comp2samp(da, da_ref, dim='time')
    plt.rcParams['hatch.color'] = 'C2'
    cs = ds['pvalue'].where(ds['pvalue']<0.05).pipe(roll_lon).plot.contourf(ax=ax, colors='none', hatches=['....'], add_colorbar=False)
    #for collection in cs.collections:
    #    collection.set_facecolor('C2')
    plt.rcParams['hatch.color'] = 'gray'

    
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    if daname == 't_surf':
        vmax = 1
    elif daname == 'slp':
        vmax = None #3 
    elif daname == 'precip':
        vmax = 2 
    else:
        vmax = None
    if daname == 'precip':
        cmap = 'BrBG'
    else:
        cmap = None
    for da,label in zip(das[1::2], labels[1::2]):
        fig,ax = plt.subplots()
        year_start = da.time[0].item().year
        year_stop = da.time[-1].item().year
        wyplot(da, label, ax=ax, levels=21, vmax=vmax, cmap=cmap, extend='both')
        plt.sca(ax)
        mapplot()

        ax.set_title(f"{label} {season} {daname} anomaly: {year_start}-{year_stop}")

        #savefig
        if 'savefig' in sys.argv or 's' in sys.argv:
            figname = __file__.replace('.py', f'__{daname}_{label}_{season}.png')
            if 'overwritefig' in sys.argv or 'o' in sys.argv:
                wysavefig(figname, overwritefig=True)
            else:
                wysavefig(figname)
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv:
        pass
    else:
        plt.show()