#!/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.comp1samp import comp1samp
from misc.comp2samp import comp2samp
from misc.landmask import whereocean
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
model = 'AM2.5'
#daname = 'netrad_toa'
#daname = 'netrad_toa_clr'
#daname = 'swdn_sfc'
#daname = 'tot_cld_amt'
#daname = get_kws_from_argv('daname', daname)
danames = """
netrad_toa
olr
swup_toa
swdn_sfc
swup_sfc
lwdn_sfc
lwup_sfc
shflx
evap
precip
""".split()
func = lambda da: da.rename(grid_xt='lon', grid_yt='lat').pipe(whereocean).geo.fldmean() #  da.geo.fldmean()
funcname = 'oceanmean'#'glbmean'
years = range(101,131)

das_danames = dict()
for daname in danames:
    print(daname)
    das_cases = dict()
    #ctl
    expname = 'CTL1990s_tigercpu_intelmpi_18_540PE'
    da = get_modelout_data(daname=daname, model=model, expname=expname, years=years, func=func, funcname=funcname)
    if daname == 'evap':
        da = (da*2.5e6).assign_attrs(units='W/m^2')
    elif daname == 'precip':
        da = (da*2.5e6/24/3600).assign_attrs(units='W/m^2')
    das_cases['ctl'] = da.groupby('time.year').mean('time')
    units = da.attrs['units']
    #perturb
    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]
    for case,expname in zip(cases,expnames):
        print(case)
        da = update_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)
        if daname == 'evap':
            da = (da*2.5e6).assign_attrs(units='W/m^2')
        elif daname == 'precip':
            da = (da*2.5e6/24/3600).assign_attrs(units='W/m^2')
        das_cases[case] = da.groupby('time.year').mean('time')
    das_danames[daname] = das_cases

#compound variables
#
daname = 'netflx_sfc'
das_danames[daname] = dict()
for case in ['ctl',]+cases:
    das_danames[daname][case] = das_danames['swdn_sfc'][case] \
        - das_danames['swup_sfc'][case] \
        + das_danames['lwdn_sfc'][case] \
        - das_danames['lwup_sfc'][case] \
        - das_danames['evap'][case] \
        - das_danames['shflx'][case]
danames.append(daname)
#
daname = 'netrad_sfc'
das_danames[daname] = dict()
for case in ['ctl',]+cases:
    das_danames[daname][case] = das_danames['swdn_sfc'][case] \
        - das_danames['swup_sfc'][case] \
        + das_danames['lwdn_sfc'][case] \
        - das_danames['lwup_sfc'][case]
danames.append(daname)
#
daname = 'netrad_heating'
das_danames[daname] = dict()
for case in ['ctl',]+cases:
    das_danames[daname][case] = das_danames['netrad_toa'][case] - das_danames['netrad_sfc'][case]
danames.append(daname)
#
daname = 'netsw_sfc'
das_danames[daname] = dict()
for case in ['ctl',]+cases:
    das_danames[daname][case] = das_danames['swdn_sfc'][case] - das_danames['swup_sfc'][case]
danames.append(daname)
#
daname = 'netlw_sfc'
das_danames[daname] = dict()
for case in ['ctl',]+cases:
    das_danames[daname][case] = das_danames['lwdn_sfc'][case] - das_danames['lwup_sfc'][case]
danames.append(daname)

#calculate mean diff and ci
dss_danames = []
for daname in danames:
    dss_cases = []
    for case in cases:
        ds = comp2samp(das_danames[daname][case], das_danames[daname]['ctl'], dim='year')
        dss_cases.append(ds)
    ds = xr.concat(dss_cases, dim=pd.Index(cases, name='case'))
    dss_danames.append(ds)
ds = xr.concat(dss_danames, dim=pd.Index(danames, name='daname'))


def wyplot(cases=None, **kws):
    if cases is None:
        cases = ['2xCO2', 'solar2pctplus']
    danames_show = ['netrad_toa', 'netflx_sfc', 'netrad_sfc', 'netrad_heating', 'evap', 'shflx', 'netsw_sfc', 'netlw_sfc', 'precip']
    df = ds.sel(case=cases).sel(daname=danames_show)['x1m_minus_x2m'].to_pandas()
    df_err = ds.sel(case=cases).sel(daname=danames_show)['err'].to_pandas()
    df.plot.bar(rot=-45, yerr=df_err, capsize=3, table=df.T.round(1))
    
    ax = plt.gca()
    #ax.set_xticklabels('')
    #ax.get_xaxis().set_visible(False)
    #ax.xaxis.tick_top()
    ax.set_xlabel('')
    ax.set_ylabel('W m$^{-2}$')
    ax.axhline(0, color='gray', ls='-', lw=1)
    
    
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    plt.close('all')
    cases = ['2xCO2', 'solar2pctplus']
    #cases = ['2xCO2', 'p5xCO2']
    #cases = ['solar2pctminus', 'solar2pctplus']
    wyplot(cases=cases)

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