#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Nov  3 18:53:46 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
import geoxarray
from misc.modelout import get_modelout_data, update_modelout_data
from misc import get_kws_from_argv
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
model = 'AM2.5'
#daname = 'netrad_toa_clr'
#daname = 'netrad_toa'
#daname = 'swup_toa_clr'
daname = 'olr_clr'
#daname = 'swdn_sfc'
#daname = 'tot_cld_amt'
daname = get_kws_from_argv('daname', daname)
expname = 'CTL1990s_tigercpu_intelmpi_18_540PE'
func = None
funcname = 'wy'
years = range(101,131)
das = dict()

#ctl
da = get_modelout_data(daname=daname, model=model, expname=expname, years=years, func=func, funcname=funcname)
das['ctl'] = da
units = da.attrs['units']

expnames = """
CTL1990s_stratBC1e-09kgkg_tigercpu_intelmpi_18_540PE
CTL1990s_stratBC2e-09kgkg_tigercpu_intelmpi_18_540PE
CTL1990s_stratBC4e-09kgkg_tigercpu_intelmpi_18_540PE
CTL1990s_stratBC8e-09kgkg_tigercpu_intelmpi_18_540PE
CTL1990s_stratBC8e-09kgkg_5-10hPa_tigercpu_intelmpi_18_540PE
CTL1990s_stratBC8e-09kgkg_10-20hPa_tigercpu_intelmpi_18_540PE
""".split()[3:4]
years = range(103,105+1)
cases = [expname.split('CTL1990s_')[1].split('_tigercpu_')[0] for expname in expnames]
for case,expname in zip(cases,expnames):
    #da = update_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)
    da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, years=years)
    das[case] = da
case = cases[0]

def wyplot(case, **kws):
    daa = das[case].mean('time') - das['ctl'].mean('time')
    #daa.plot(marker='o', fillstyle='none', label=case, **kws)
    daa.assign_attrs(units=units, long_name=daname+' anom').plot.contourf(levels=21, **kws)

    ax.set_title(f'{daname} anom, {case}')
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig,ax = plt.subplots()
    wyplot(case)
    mapplot()

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