#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sat Feb 17 15:08:05 EST 2024
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
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
#from misc import get_kws_from_argv
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
yearspan = slice(1981,2010)
nyears = yearspan.stop - yearspan.start + 1
windmax2min = (17, 33, 50)
tc_types = ('TC', 'HU', 'MH')
datafile = __file__.replace('.py', '.nc')
datafile_sem = datafile.replace('.nc', '__sem.nc') #standard error of the mean
if os.path.exists(datafile):
    ds = xr.open_dataset(datafile)
    ds_sem = xr.open_dataset(datafile_sem)
    print('[loaded]:', datafile)
    print('[loaded]:', datafile_sem)
else:
    dss = []
    dss_sem = []
    for ww,tc_type in zip(windmax2min, tc_types):
        ifile = f'IBTrACS.ALL.v04r00.2022-01-25.counts.TS{ww}.1980-2021.yearly.nc'
        ds = xr.open_dataset(ifile).sel(year=yearspan).mean('year')
        ds_sem = xr.open_dataset(ifile).sel(year=yearspan).std('year')/np.sqrt(nyears)
        dss.append(ds)
        dss_sem.append(ds_sem)
    ds = xr.concat(dss, dim=pd.Index(tc_types, name='cat'))
    ds_sem = xr.concat(dss_sem, dim=pd.Index(tc_types, name='cat'))
    ds.to_netcdf(datafile)
    print('[saved]:', datafile)
    ds_sem.to_netcdf(datafile_sem)
    print('[saved]:', datafile_sem)
df = ds.to_pandas().iloc[:,2:]
df_sem = ds_sem.to_pandas().iloc[:,2:]
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig,ax = plt.subplots()
    df.T.plot.bar(ax=ax, rot=0, yerr=2*df_sem.T, table=np.round(df,2))

    ax.set_xticklabels('')
    ax.set_title(f'IBTrACS annual mean # of storms, {yearspan.start}-{yearspan.stop}')
    ax.set_ylabel('#')
    
    #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 or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()