#!/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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
model = 'AM2.5'
daname = 'temp'
#daname = 'cld_amt'
expname = 'CTL1990s_tigercpu_intelmpi_18_540PE'
func = lambda da: da.geo.fldmean()
funcname = 'glbmean'
years_ctl = range(101,131)
das = dict()

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

"""
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
CTL1990s_solar1pctplus_tigercpu_intelmpi_18_540PE
"""
expnames = """
CTL1990s_p5xCO2_tigercpu_intelmpi_18_540PE
CTL1990s_2xCO2_tigercpu_intelmpi_18_540PE
CTL1990s_solar2pctminus_tigercpu_intelmpi_18_540PE
CTL1990s_solar2pctplus_tigercpu_intelmpi_18_540PE
""".split()
cases = [expname.split('CTL1990s_')[1].split('_tigercpu_')[0] for expname in expnames]
years = range(111,131)
for case,expname in zip(cases,expnames):
    da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, years=years)
    das[case] = da

def wyplot(case, flip=False, **kws):
    daa = das[case].groupby('time.month') - das['ctl'].groupby('time.month').mean('time')
    #daa = daa.isel(time=slice(-12,None)).mean('time') # value of last year
    daa = daa.mean('time') # 
    #daa.plot(marker='o', fillstyle='none', label=case, **kws)
    linestyle = '-'
    if flip:
        if 'p5x' in case or 'minus' in case:
            daa = -daa
            linestyle = '--'
    if daname in ('temp',):
        #daa.assign_attrs(units=units, long_name=daname+' anom').plot(label=case, y='pfull', yincrease=False, yscale='log', marker='o', fillstyle='none', **kws)
        daa.assign_attrs(units=units, long_name=daname+' anom').plot(label=case, y='pfull', yincrease=False, yscale='log', ls=linestyle, **kws)
    else:
        daa.assign_attrs(units=units, long_name=daname+' anom').plot(label=case, y='pfull', yincrease=False, ls=linestyle, **kws)

    ax.set_title(f'{daname} anom, {model} {years[0]}-{years[-1]}')
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    flip = True if 'flip' in sys.argv else False
    fig,ax = plt.subplots()
    for case in cases:
        wyplot(case, flip=flip)
    ax.legend(ncol=1)
    #for plevel in (5,10,20,100):
    for plevel in (100,):
        ax.axhline(plevel, color='gray', ls='--')

    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__{daname}_{years[0]}-{years[-1]}.png')
        if flip: figname = figname.replace('.png', '_flip.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()