#!/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']
#obs
ifile = '/tigress/wenchang/data/ibtracs/v04r00/analysis/v2/IBTrACS.ALL.v04r00.2022-01-25.counts.1980-2021.yearly.nc'
da = xr.open_dataset(ifile).pipe(ds2da)
#print(da)
da1 = da.sel(year=slice(2001,2020))
da2 = da.sel(year=slice(1981,1990))
result = comp2samp(da1, da2, dim='year')

ds = xr.Dataset({
    'obs2010s-1980s': result['x1m_minus_x2m']/result['x2m']*100,
    'err_obs': result['err']/result['x2m']*100,
    })

#amip tc
ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.10ens.1871-2021.yearly.nc'
da = xr.open_dataset(ifile).pipe(ds2da)
#print(da)
da1 = da.sel(year=slice(2001,2020)).stack(s=['en', 'year'])
da2 = da.sel(year=slice(1971,1990)).stack(s=['en', 'year'])
result = comp2samp(da1, da2, dim='s')

"""
ds = xr.Dataset({
    'AMIP2010s-1980s': result['x1m_minus_x2m']/result['x2m']*100,
    'err_amip': result['err']/result['x2m']*100,
    })
"""
ds['AMIP2010s-1980s'] = result['x1m_minus_x2m']/result['x2m']*100
ds['err_amip'] = result['err']/result['x2m']*100

#time slice tc
ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL2010s_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.0101-0150.yearly.nc'
da1 = xr.open_dataset(ifile).pipe(ds2da)
ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL1980s_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.0101-0150.yearly.nc'
da2 = xr.open_dataset(ifile).pipe(ds2da)
result = comp2samp(da1, da2, dim='year')
ds['TimeSlice2010s-1980s'] = result['x1m_minus_x2m']/result['x2m']*100
ds['err_tslice'] = result['err']/result['x2m']*100

#time slice tc, ice1980
ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL2010s_ice1980_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.0101-0150.yearly.nc'
da1 = xr.open_dataset(ifile).pipe(ds2da)
ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL1980s_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.0101-0150.yearly.nc'
da2 = xr.open_dataset(ifile).pipe(ds2da)
result = comp2samp(da1, da2, dim='year')
ds['TimeSlice2010s-1980s_fixIce'] = result['x1m_minus_x2m']/result['x2m']*100
ds['err_tslice_fixIce'] = result['err']/result['x2m']*100

#time slice tc, ice10 - ice1980
ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL2010s_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.0101-0150.yearly.nc'
da1 = xr.open_dataset(ifile).pipe(ds2da)
ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL2010s_ice1980_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.0101-0150.yearly.nc'
da2 = xr.open_dataset(ifile).pipe(ds2da)
result = comp2samp(da1, da2, dim='year')
ds['TimeSlice2010s-2010s_onlyIce'] = result['x1m_minus_x2m']/result['x2m']*100
ds['err_tslice_onlyIce'] = result['err']/result['x2m']*100

ds = ds.isel(basin=slice(None,-1))
df = ds.to_pandas() 
 
def wyplot(frame=None, include_obs=False, **kws):
    y=['obs2010s-1980s', 'AMIP2010s-1980s', 'TimeSlice2010s-1980s', 'TimeSlice2010s-1980s_fixIce', 'TimeSlice2010s-2010s_onlyIce']
    yerr=[df['err_obs'], df['err_amip'], df['err_tslice'], df['err_tslice_fixIce'], df['err_tslice_onlyIce']]
    if not include_obs:
        y = y[1:]
        yerr = yerr[1:]
    if frame is not None:
        y = y[0:frame]
        yerr = yerr[0:frame]
    #df.plot.bar(y=y, yerr=yerr, capsize=3, rot=0, **kws)
    columns_new = {k:' ' for k in y}
    table = df[y].T.round(1)#.rename(columns_new)
    #columns_new = {k:' ' for k in table.columns.values}
    index_new = {k:' ' for k in table.index.values}
    table = table.rename(index=index_new)
    df.plot.bar(y=y, yerr=yerr, capsize=3, rot=0, table=table, **kws)

    ax = plt.gca()
    ax.axhline(0, color='gray', lw=1)
    ax.set_ylim(-30,50)
    ax.set_ylabel('%')
    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 = True if 'frames' in sys.argv else False
    if frames:
        plt.close('all')
        nframe = df.columns.size//2
        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()