#!/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
month0 = get_kws_from_argv('month0', 'Jan')
amipyears = range(1997, 1998+1)
if month0 == 'May':
    amipyears = list(range(1997, 1998+1)) + [2026,]
#cases = ['CTL8100', 'CTL9120', 'CTL0120'] + list(amipyears)
cases = ['CTL8100', 'CTL0120'] + list(amipyears)
das_amip = []
das_hindcast = []

#CTL8100
ifile = '/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/CTL8100/TC/tc.counts.NA.nc'
da = xr.open_dataarray(ifile)
das_amip.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')
#new code
ifile = '/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/CTL9120/TC/tc.counts.NA.nc'
da = xr.open_dataarray(ifile)
das_amip.append(da)
"""

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

#hindcast
for amipyear in amipyears:
    """
    #old code
    ifile = f'/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/Hindcast{amipyear}01/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')
    """
    #hindcast
    ifile = f'/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/Hindcast{amipyear}01/TC/tc.counts.NA.nc'
    if month0 == 'May':
        ifile = ifile.replace('01', '05')
        if amipyear == 2026:
            ifile = ifile.replace('Hind', 'Fore')
    da = xr.open_dataarray(ifile)
    das_hindcast.append(da)
    #amip
    if amipyear == 2026:
        continue
    ifile = f'/projects/GEOCLIM/wenchang/MODEL_OUT/ACE2/AMIP{amipyear}/TC/tc.counts.NA.nc'
    da = xr.open_dataarray(ifile)
    das_amip.append(da)

da_amip = xr.concat(das_amip, dim=pd.Index(cases[:-1], name='case'))
da_amip = da_amip.isel(year=slice(0,100)) #sel the first 100 years to have the same size as hindcast
print(da.mean('year'))
print(da.median('year'))
da_hindcast = xr.concat(das_hindcast, dim=pd.Index(cases[2:], name='case'))
print(da.mean('year'))
print(da.median('year'))
df_amip = da_amip.rename(year='ens').stack(s=['case', 'ens']).to_dataframe()
df_amip['hindcast'] = False
df_hindcast = da_hindcast.rename(year='ens').stack(s=['case', 'ens']).to_dataframe()
df_hindcast['hindcast'] = True
df = pd.concat([df_amip, df_hindcast])
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='hindcast', gap=0.05, linewidth=1)

    ax = plt.gca()
    if month0 == 'Jan':
        ax.set_title('ACE2 100-member NA NTC HindcastJanIC')
    elif month0 == 'May':
        ax.set_title('ACE2 100-member NA NTC HindcastMayIC')
    ax.set_ylabel('#')
    ax.set_xlabel('')
    #ax.set_xticklabels(ax.get_xticklabels(), rotation=30, ha='right')

    ax.legend(handles=ax.legend_.legend_handles, labels=['AMIP', 'Hindcast',])
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        figname = figname.replace('.png', f'__{month0}.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()