#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Nov 21 10:03:26 EST 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
from misc.comp2samp import comp2samp
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
def ds2da(ds):
    """convert tc counts ds to da"""
    if 'global' in ds:
        basins = list(ds.data_vars)
        return xr.concat([ds[basin] for basin in basins], dim=pd.Index(basins, name='basin'))
    else:
        return ds['tc_count']

dss = []
cases = []
#p5xCO2
ifile = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_p5xCO2_tigercpu_intelmpi_18_540PE/netcdf/tc_counts.TS17.0101-0150.yearly.nc'
da1 = xr.open_dataset(ifile).pipe(ds2da)
ifile = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_tigercpu_intelmpi_18_540PE/netcdf/tc_counts.TS17.0101-0200.yearly.nc'
da2 = xr.open_dataset(ifile).pipe(ds2da)
ds = comp2samp(da1, da2, dim='year')
dss.append(ds)
cases.append('0.5xCO2')

#2xCO2
ifile = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_2xCO2_tigercpu_intelmpi_18_540PE/netcdf/tc_counts.TS17.0101-0150.yearly.nc'
da1 = xr.open_dataset(ifile).pipe(ds2da)
ifile = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_tigercpu_intelmpi_18_540PE/netcdf/tc_counts.TS17.0101-0200.yearly.nc'
da2 = xr.open_dataset(ifile).pipe(ds2da)
ds = comp2samp(da1, da2, dim='year')
dss.append(ds)
cases.append('2xCO2')

#-2% solar
ifile = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_solar2pctminus_tigercpu_intelmpi_18_540PE/netcdf/tc_counts.TS17.0101-0150.yearly.nc'
da1 = xr.open_dataset(ifile).pipe(ds2da)
ifile = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_tigercpu_intelmpi_18_540PE/netcdf/tc_counts.TS17.0101-0200.yearly.nc'
da2 = xr.open_dataset(ifile).pipe(ds2da)
ds = comp2samp(da1, da2, dim='year')
dss.append(ds)
cases.append('$-$2%_solar')

#+2% solar
ifile = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_solar2pctplus_tigercpu_intelmpi_18_540PE/netcdf/tc_counts.TS17.0101-0150.yearly.nc'
da1 = xr.open_dataset(ifile).pipe(ds2da)
ifile = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_tigercpu_intelmpi_18_540PE/netcdf/tc_counts.TS17.0101-0200.yearly.nc'
da2 = xr.open_dataset(ifile).pipe(ds2da)
ds = comp2samp(da1, da2, dim='year')
dss.append(ds)
cases.append('$+$2%_solar')

ds = xr.concat(dss, dim=pd.Index(cases, name='case'))
ds['delta_pct'] = (ds['x1m_minus_x2m']/ds['x2m']*100).assign_attrs(units='%')
ds['err_pct'] = (ds['err']/ds['x2m']*100).assign_attrs(units='%')


#remove the SA basin
ds = ds.isel(basin=slice(None,-1))
 
def wyplot(frame=None, **kws):
    #to pandas dataframe
    if frame is None:
        df = ds['delta_pct'].to_pandas().T
        df_err = ds['err_pct'].to_pandas().T 
    else:
        df = ds['delta_pct'].isel(None, frame).to_pandas().T
        df_err = ds['err_pct'].isel(None, frame).to_pandas().T 

    df.plot.bar(yerr=df_err, capsize=3, rot=-45, table=df.T.round(1), **kws)

    ax = plt.gca()
    ax.axhline(0, color='gray', lw=1)
    #plt.ylim(-30,50)
    ax.set_ylabel('AM2.5 TC change [%]')
    ax.set_xlabel('')
    ax.set_xticklabels('')
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from misc import get_kws_from_argv
    #fig,ax = plt.subplots()
    #ax.bar('basin', ['AMIP2010s-1980s', 'TimeSlice2010s-1980s'], yerr=['err_amip', 'err_tslice'], capsize=3, data=ds)
    frames = False #True if 'frames' in sys.argv else False
    if frames:
        plt.close('all')
        nframe = ds.case.size
        for ii in range(1,nframe+1):
            wyplot(frame=ii)
            figname = __file__.replace('.py', f'__frame{ii}.png')
            wysavefig(figname, overwritefig=True)
    else:
        wyplot()
        #savefig
        if 'savefig' in sys.argv or 's' in sys.argv:
            figname = __file__.replace('.py', f'.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()