#!/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
from misc.landmask import whereland
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
model = 'AM2.5'
#daname = 'netrad_toa_clr'
#daname = 'netrad_toa'
#daname = 'swdn_sfc'
#daname = 'swup_sfc'
#daname = 'swup_toa'
#daname = 'lwup_sfc'
#daname = 'lwdn_sfc'
#daname = 'shflx'
daname = 't_surf'
#daname = 'tot_cld_amt'
daname = get_kws_from_argv('daname', daname)
expname = 'CTL1990s_tigercpu_intelmpi_18_540PE'
func = lambda da: da.rename(grid_xt='lon', grid_yt='lat').pipe(whereland).geo.fldmean()
funcname = 'landmean'
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_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
CTL1990s_stratBC8e-09kgkg_15-25hPa_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):
    stopyear = kws.pop('stopyear', None)
    ax = kws.pop('ax', plt.gca())
    daa = das[case].groupby('time.month') - das['ctl'].groupby('time.month').mean('time')
    #daa.plot(marker='o', fillstyle='none', label=case, **kws)
    if stopyear is not None:
        #f = lambda da: da.rolling(time=12, center=False, min_periods=1).mean().sel(time=slice(None, '0110'))
        f = lambda da: da.sel(time=slice(None, f'{stopyear}'))
    else:
        f = lambda da: da
    daa.pipe(f).assign_attrs(units=units, long_name=daname+' anom').plot(label=case+f', {daa.mean().item():.3g}', **kws)
    ax.axhline(daa.mean(), ls='--', **kws)


    ax.set_title(f'{daname} {funcname} anom, {model}')
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from misc import get_kws_from_argv
    stopyear = get_kws_from_argv('stopyear', None)
    fig,ax = plt.subplots()
    for ii,case in enumerate(cases):
        wyplot(case, color=f'C{ii}', stopyear=stopyear)
    ax.legend(ncol=1)

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