#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sun Feb  5 18:05:53 EST 2023
if __name__ == '__main__':
    import sys,os
    from misc.timer import Timer
    tt = Timer(f'[{os.getcwd()}] 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 xtc import tc_count
from misc.comp2samp import comp2samp
from misc import get_kws_from_argv
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
case = 'FLORktc2' #'AM2.5plus2K'
case = get_kws_from_argv('case', case)
basin = 'global'
ws = 17
ws = get_kws_from_argv('ws', ws)
ws = int(ws)
if case.startswith('FLOR'):
    model = 'FLOR'
elif case.startswith('AM2.5'):
    model = 'AM2.5'
    if case.startswith('AM2.5C360'):
        model = 'AM2.5C360'
elif case.startswith('HIRAM'):
    model = 'HIRAM'

#ctl
label_ctl = f'{basin}{ws}{model}'
ifile_ctl = '/tigress/wenchang/analysis/TC/CTL1990_v201905_tigercpu_intelmpi_18_576PE/netcdf/tc_tracks.TS.0001-1000.nc'
#label 
label = f'{basin}{ws}{case}'
ifile = '/tigress/wenchang/analysis/TC/FLORktc2/CTL1990_v201905_killtc13-13-15_tigercpu_intelmpi_18_576PE/netcdf/tc_tracks.TS.0101-0300.nc'
#
if case == 'FLORktc2':
    pass #the default
elif case == 'AM2.5plus2K':
    ifile_ctl = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0200.nc'
    ifile = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_plus2K_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'AM2.52xCO2':
    ifile_ctl = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0200.nc'
    ifile = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_2xCO2_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'AM2.5plus2K2xCO2':
    ifile_ctl = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0200.nc'
    ifile = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_plus2K_2xCO2_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'AM2.5ktc2':
    ifile_ctl = '/tigress/wenchang/analysis/TC/AM2.5/CTL1990s_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0200.nc'
    ifile = '/tigress/wenchang/analysis/TC/AM2.5ktc2/CTL1990s_killtc13-13-15_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'AM2.5ktc2plus2K':
    label_ctl = f'{basin}{ws}{model}ktc2'
    ifile_ctl = '/tigress/wenchang/analysis/TC/AM2.5ktc2/CTL1990s_killtc13-13-15_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0150.nc'
    label = f'{basin}{ws}{case}'
    ifile = '/tigress/wenchang/analysis/TC/AM2.5ktc2/CTL1990s_killtc13-13-15_plus2K_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'HIRAMplus2K':
    ifile_ctl = '/tigress/wenchang/analysis/TC/HIRAM/CTL1990s_v201910_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0300.nc'
    ifile = '/tigress/wenchang/analysis/TC/HIRAM/CTL1990s_v201910_plus2K_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'HIRAM2xCO2':
    ifile_ctl = '/tigress/wenchang/analysis/TC/HIRAM/CTL1990s_v201910_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0300.nc'
    ifile = '/tigress/wenchang/analysis/TC/HIRAM/CTL1990s_v201910_2xCO2_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'HIRAMplus2K2xCO2':
    ifile_ctl = '/tigress/wenchang/analysis/TC/HIRAM/CTL1990s_v201910_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0300.nc'
    ifile = '/tigress/wenchang/analysis/TC/HIRAM/CTL1990s_v201910_plus2K_2xCO2_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'HIRAMktc2':
    ifile_ctl = '/tigress/wenchang/analysis/TC/HIRAM/CTL1990s_v201910_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0300.nc'
    ifile = '/tigress/wenchang/analysis/TC/HIRAMktc2/CTL1990s_v201910_killtc13-13-15_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'HIRAMktc2plus2K':
    label_ctl = f'{basin}{ws}{model}ktc2'
    ifile_ctl = '/tigress/wenchang/analysis/TC/HIRAMktc2/CTL1990s_v201910_killtc13-13-15_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0150.nc'
    label = f'{basin}{ws}{case}'
    ifile = '/tigress/wenchang/analysis/TC/HIRAMktc2/CTL1990s_v201910_killtc13-13-15_plus2K_tigercpu_intelmpi_18_540PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'AM2.5C360plus2K':
    ifile_ctl = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL1990s_tigercpu_intelmpi_18_1080PE/netcdf/tc_tracks.TS.0101-0200.nc'
    ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL1990s_plus2K_tigercpu_intelmpi_18_1080PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'AM2.5C3602xCO2':
    ifile_ctl = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL1990s_tigercpu_intelmpi_18_1080PE/netcdf/tc_tracks.TS.0101-0200.nc'
    ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL1990s_2xCO2_tigercpu_intelmpi_18_1080PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'AM2.5C360plus2K2xCO2':
    ifile_ctl = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL1990s_tigercpu_intelmpi_18_1080PE/netcdf/tc_tracks.TS.0101-0200.nc'
    ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL1990s_plus2K_2xCO2_tigercpu_intelmpi_18_1080PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'AM2.5C360ktc2':
    ifile_ctl = '/tigress/wenchang/analysis/TC/AM2.5C360/CTL1990s_tigercpu_intelmpi_18_1080PE/netcdf/tc_tracks.TS.0101-0200.nc'
    ifile = '/tigress/wenchang/analysis/TC/AM2.5C360ktc2/CTL1990s_killtc13-13-15_tigercpu_intelmpi_18_1080PE/netcdf/tc_tracks.TS.0101-0150.nc'
elif case == 'AM2.5C360ktc2plus2K':
    label_ctl = f'{basin}{ws}{model}ktc2'
    ifile_ctl = '/tigress/wenchang/analysis/TC/AM2.5C360ktc2/CTL1990s_killtc13-13-15_tigercpu_intelmpi_18_1080PE/netcdf/tc_tracks.TS.0101-0150.nc'
    label = f'{basin}{ws}{case}'
    ifile = '/tigress/wenchang/analysis/TC/AM2.5C360ktc2/CTL1990s_killtc13-13-15_plus2K_tigercpu_intelmpi_18_1080PE/netcdf/tc_tracks.TS.0101-0150.nc'
else:
    print('please select the right case')
    sys.exit()

ofile = __file__.replace('.py', f'_{label_ctl}_{label}.nc')
if os.path.exists(ofile):
    ds = xr.open_dataset(ofile)
    print('[loaded]:', ofile)
else:
    #get yearly TCs
    da_ctl = xr.open_dataset(ifile_ctl) \
        .pipe(tc_count, basin=basin, ws=ws) \
        .groupby('time.year').sum('time')
    da = xr.open_dataset(ifile) \
        .pipe(tc_count, basin=basin, ws=ws) \
        .groupby('time.year').sum('time')
    #anomaly
    ds = comp2samp(da, da_ctl)
    #save
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)

#convert to dataframe
df = ds[['x2m', 'x1m']].rename(x2m=label_ctl, x1m=label).to_pandas()
df_s = ds[['s2', 's1']].rename(s2=label_ctl, s1=label).to_pandas()
df_n = ds[['n2', 'n1']].rename(n2=label_ctl, n1=label).to_pandas()
df_e = df_s/np.sqrt(df_n)

 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    colors = [f'C{ii}' for ii in range(df.size)]
    df.plot.bar(color=colors, rot=0, yerr=df_e, capsize=3)
    ax = plt.gca()
    for ii in range(df.size):
        if ii<df.size-1:
            ax.text(ii, df[ii], f'{df[ii]:.1f}\n', ha='center', va='bottom')
        else:
            ax.text(ii, df[ii], f'{df[ii]:.1f} ({(df[ii]-df[0])/df[0]:.1%})\n', ha='center', va='bottom')
            #ax.text(ii, df[ii], f'{(df[ii]-df[0])/df[0]:.1%}', ha='center', va='bottom')
    ax.set_ylabel('#')
            
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'_{label_ctl}_{label}.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 or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()