#!/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
from misc.seasons import sel_season
from misc import get_kws_from_argv
import xfilter
#nwindow, dimlp = 9, 'year'
nwindow, dimlp = 1, '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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
daname = 't_ref'
#func = lambda x: x.load().geo.fldmean()
#funcname = 'glbmean'
funcname = 'none'
func = funcs[funcname]

das = []
labels = []


model = 'ESM4'
#CTL1850, t_ref
label = 'ESM4_ctl_1850 t_ref'
expname = 'CTL1850_tiger3_intel24ifort_openmpi_2296PE'
da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, years=range(1,101))
labels.append(label)
das.append(da)
#CTL1850, t_surf
daname = 't_surf'
label = 'ESM4_ctl_1850 t_surf'
expname = 'CTL1850_tiger3_intel24ifort_openmpi_2296PE'
da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, years=range(1,101))
labels.append(label)
das.append(da)

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    season = get_kws_from_argv('season', 'annual')
    fig,ax = plt.subplots()
    with xr.set_options(keep_attrs=True):
        daa = das[1].load().pipe(sel_season, season).mean('time') - das[0].load().pipe(sel_season, season).mean('time')
        long_name = f'{season} t_surf $-$ t_ref'
        daa = daa.assign_attrs(long_name=long_name)
        daa.plot.contourf(ax=ax, levels=23, vmax=11)
    mapplot()
    

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__{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()