#!/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
cases = ['CTL8100', 'CTL9120', 'CTL0120']
das = []

#CTL8100
ifile = '/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/CTL8100/TC/tc.counts.NA.nc'
da = xr.open_dataarray(ifile)
print(da)
das.append(da)

#CTL9120
"""
#old code
ifile = '/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/CTL9120/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 = '/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/CTL9120/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)
print(da)
das.append(da)

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

#plot
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    df.pipe(sns.violinplot,x='case', y='tc_counts', cut=0, split=True, inner='quart', hue='1000_members', gap=0.05, linewidth=1)

    ax = plt.gca()
    ax.set_title('ACE2 large ensemble NA NTC')
    ax.set_ylabel('#')
    ax.set_xlabel('')
    #ax.set_xticklabels(ax.get_xticklabels(), rotation=30, ha='right')

    ax.legend(handles=ax.legend_.legend_handles, labels=['100-member', '1000-member',])
    
    #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()