#!/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 = 'swdn_sfc'
#daname = 'tot_cld_amt'
daname = get_kws_from_argv('daname', daname)
expname = 'CTL1990s_tigercpu_intelmpi_18_540PE'
func = lambda da: da.geo.fldmean()
funcname = 'glbmean'
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_old = """
CTL1990s_stratBC2e-09kgkg_tigercpu_intelmpi_18_540PE
CTL1990s_stratBC4e-09kgkg_tigercpu_intelmpi_18_540PE
"""
expnames = """
CTL1990s_stratBC1e-09kgkg_tigercpu_intelmpi_18_540PE
CTL1990s_stratBC8e-09kgkg_tigercpu_intelmpi_18_540PE
CTL1990s_stratBC8e-09kgkg_5-10hPa_tigercpu_intelmpi_18_540PE
CTL1990s_stratBC2e-08kgkg_5-10hPa_tigercpu_intelmpi_18_540PE
CTL1990s_stratBC8e-09kgkg_10-20hPa_tigercpu_intelmpi_18_540PE
CTL1990s_stratBC8e-09kgkg_45-55hPa_tigercpu_intelmpi_18_540PE
""".split()
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)
    das[case] = da

def wyplot(case, **kws):
    daa = das[case].groupby('time.month') - das['ctl'].groupby('time.month').mean('time')
    #daa.plot(marker='o', fillstyle='none', label=case, **kws)
    #N = daa.sel(time=slice('0106', '0110')).mean('time').item()
    N = daa.isel(time=slice(-36,None)).mean('time').item() #last three years avergae
    daa.assign_attrs(units=units, long_name=daname+' anom').plot(label=case.split('BC')[1]+f', {N:.2g}', **kws)

    ax.set_title(f'{daname} anom, stratBC - CTL')
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig,ax = plt.subplots()
    for case in cases:
        wyplot(case)
    ax.legend(ncol=2)

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