#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Oct 30 15:12:44 EDT 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
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
from xtc import tracks_in_basin
from xtc import trackplot
import xtc
from geoplots import mapplot
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
basin = 'EP'
#ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE/netcdf/tc_tracks.TS.10ens.1871-2021.nc'
ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_rmWarming_tigercpu_intelmpi_18_1080P/netcdf/tc_tracks.TS.3ens.1871-2020.nc'
ds = xr.open_dataset(ifile).load()
ds_ep = ds.where( ds.pipe(tracks_in_basin, basin) )

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig,ax = plt.subplots()
    (ds_ep.windmax.max('stage')>50).sum('storm').sum('en').plot(label='Cat3-5', ax=ax)
    (ds_ep.windmax.max('stage')>58).sum('storm').sum('en').plot(label='Cat4-5', ax=ax)
    (ds_ep.windmax.max('stage')>70).sum('storm').sum('en').plot(label='Cat5', ax=ax)

    ax.legend()
    n_ens = ds.en.size
    ax.set_ylabel(f'{basin} HU # per {n_ens} ens per year')
    
    #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)

    fig,ax = plt.subplots()
    L = ds_ep.windmax.max('stage')>70
    ds_ep.where(L).sel(year=1897, en=8).pipe(xtc.trackplot, lw=2, alpha=1,  ax=ax)
    ds_ep.where(L).sel(year=1897, en=8).pipe(xtc.trackplot, ezplot=True, ls='none', marker='o', fillstyle='none', color='gray', ax=ax)
    mapplot()
    ax.set_title(f'year 1897, ens 8')

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'_track_ens08year1894.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()