#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed May 13 03:06:43 PM EDT 2026
import sys
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
if __name__ == '__main__':
    try:
        from misc.timer import Timer
        tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
    except:
        pass
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 import get_kws_from_argv
from xtc.basins import tracks_in_basin
import seaborn as sns
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
#settings
#process
#amipyears = list(range(1981, 2004+1)) + list(range(2005, 2021+1))
amipyears = list(range(1981, 2021+1))
cases = list(amipyears) + ['CTL8100', 'CTL0120',] 
das = []

#hurdat2 TS
df = pd.read_csv('https://tigress-web.princeton.edu/~gvecchi/hurdat2_long_adjusted.txt', header=None, sep='\s+')
da = xr.DataArray(df.iloc[:,1].values, dims=('year',), coords=(df.iloc[:,0].astype('int'),))
da_hurdat2 = da.sel(year=slice(1981,2021))

for amipyear in amipyears:
    """
    #old code
    ifile = f'/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/AMIP{amipyear}/TC/tracks.nc'
    ds = xr.open_dataset(ifile)
    #TC basin 
    L = ds.drop('storm').isel(stage=0).pipe(tracks_in_basin, 'NA').values
    #TC count
    da = ds.lon.isel(stage=0).sel(storm=L).groupby('storm.year').count('storm')
    """
    ifile = f'/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/AMIP{amipyear}/TC/tc.counts.NA.nc'
    da = xr.open_dataarray(ifile)
    das.append(da)
#CTL8100
ifile = '/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/CTL8100/TC/tc.counts.NA.nc'
da = xr.open_dataarray(ifile)
das.append(da)
#CTL0120
ifile = '/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/CTL0120/TC/tc.counts.NA.nc'
da = xr.open_dataarray(ifile)
das.append(da)


da = xr.concat(das, dim=pd.Index(cases, name='case'))
print(da.mean('year'))
#df = da.rename(year='ens').transpose().to_pandas()
df = da.rename(year='ens').stack(s=['case', 'ens']).to_dataframe()
df['1000_members'] = True
#df['100_members'][df['case']=='CTL9120'] = True
#df['1000_members'][df['case'].isin(['CTL9120',] + list(range(1981,1996+1)))] = False
df['1000_members'][df['case'].isin(range(1981,1996+1))] = False
df100 = da.isel(year=slice(0,100)).rename(year='ens').stack(s=['case', 'ens']).to_dataframe()
df100['1000_members'] = False
df = pd.concat([df100, df])
df.reset_index(drop=True, inplace=True)
print(df)

#plot
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig,ax = plt.subplots(figsize=(14,4))
    df.pipe(sns.violinplot, x='case', y='tc_counts', cut=0, split=True, inner='quart', hue='1000_members', gap=0.2, linewidth=0.5)

    plt.axhline(da.sel(year=slice(1,100)).mean('year').isel(case=slice(0,20)).mean('case'), ls=':', label='ACE2 1981-2000 mean', color='gray') 
    plt.axhline(da.sel(year=slice(1,100)).mean('year').isel(case=slice(-23,-3)).mean('case'), ls='--', label='ACE2 2001-2020 mean', color='gray') 
    da.isel(year=slice(0,100)).mean('year').pipe(lambda x: x.where(x.case.isin(range(1981,2022)))).drop('case').plot(label='ACE2 100ens mean', color='C2', ls='-')
    da.mean('year').pipe(lambda x: x.where(x.case.isin(range(1997,2022)))).drop('case').plot(label='ACE2 1000ens mean', color='C3', ls='-')
    da.isel(year=slice(0,100)).mean('year').pipe(lambda x: x.where(x.case.isin(['CTL8100', 'CTL0120']))).drop('case').plot(color='C2', ls='-')
    da.mean('year').pipe(lambda x: x.where(x.case.isin(['CTL8100', 'CTL0120']))).drop('case').plot(color='C3', ls='-')
    plt.plot(range(da.case.size), list(da_hurdat2.values)+[np.nan, np.nan], color='k', label='HURDAT2', ls='-')

    #ax.axvspan(19.5, 39.5, color='k', alpha=0.2)
    ax.axvspan(40.5, 42.5, color='k', alpha=0.1)

    ax = plt.gca()
    ax.legend()
    ax.legend(handles=ax.legend_.legend_handles, 
        labels=['100-member', '1000-member', '1981-2000 100-ens mean', '2001-2020 100-ens mean', '100-ens mean', '1000-ens mean', 'HURDAT2'],
        ncol=2)
    ax.set_title('ACE2 large ensemble NA NTC')
    ax.set_ylabel('#')
    ax.set_xlabel('')
    ax.set_xticklabels(ax.get_xticklabels(), rotation=30, ha='right')
    
    #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)
 
    try:
        tt.check(f'**Done**')
    except:
        pass
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()