#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri May  3 15:59:59 EDT 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
import xtc
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'IBTrACS.ALL.v04r00.2022-01-25.tracksByYear.1980-2021.nc'
yearspan = slice(1980,2021)
ds = xr.open_dataset(ifile).sel(year=yearspan).load()
#sel TS
ds = ds.where(ds.windmax.max('stage')>=17)
#roll lon
ds['lon'] = ds['lon'].where(ds['lon'].isel(stage=0)>0, other=ds['lon']+360)
#interpolate windmax for NaNa values
da = ds['windmax']
#prepare data for interpolate_na: need at least two valid values along each track
da = da.where((da.stage>da.stage.min().item())|(da*0==0), other=da.min('stage')) #set the first track point, if NaN, to be the min value along the track
da = da.where((da.stage<da.stage.max().item())|(da*0==0), other=da.min('stage')) #set the last track point, if NaN, to be the min value along the track
da = da.where((da.stage>da.stage.min().item())|(da*0==0), other=15) #set the first track point, if still NaN, to be 15
da = da.where((da.stage<da.stage.max().item())|(da*0==0), other=15) #set the last track point, if still NaN, to be 15
da = da.interpolate_na(dim='stage', fill_value="extrapolate")
ds['windmax'] = da
#keep windmax data where 'lon' is available(not NaN)
ds['windmax'] = ds['windmax'].where(ds['lon']*0==0)
 
#annual cycle: TC
ifile = 'IBTrACS.ALL.v04r00.2022-01-25.counts.TS17.1980-2021.climatology.nc'
ds_cycle = xr.open_dataset(ifile).assign_coords(month=range(1,12+1))
#annual cycle: HU
ifile = 'IBTrACS.ALL.v04r00.2022-01-25.counts.TS33.1980-2021.climatology.nc'
ds_cycle_hu = xr.open_dataset(ifile).assign_coords(month=range(1,12+1))


def wyplot_basin_boundary(ax=None, **kws):
    latmax = 80
    if ax is None: fig,ax = plt.subplots()
    color = kws.pop('color', 'C2') #default color set to 'C2'
    textcolor = 'gray'
    #NA
    xy = np.array([[360-100, latmax], [360-100, 20],[360-65, 0], [360, 0], [360, latmax] ])
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    ax.text(360-100, 20, '100W', ha='right', va='bottom', color=textcolor)
    ax.text(360-65, 0, '65W', ha='left', va='bottom', color=textcolor)
    #EA
    xy = np.array([[360-160, latmax], [360-160, 0], [360-65, 0]] )
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    ax.text(360-160, 0, '160W', ha='left', va='bottom', color=textcolor)
    #WP
    xy = np.array([[105, latmax], [105, 0], [360-160, 0]] )
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    ax.text(105, 0, '105E', ha='left', va='bottom', color=textcolor)
    #NI
    xy = np.array([[30, latmax], [30, 0], [105, 0]] )
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    #SI
    xy = np.array([[30, -latmax], [30, 0], [105, 0], [105, -latmax]] )
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    #AU
    xy = np.array([[105, 0], [165, 0], [165,-latmax]] )
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    ax.text(165, -1, '165E', ha='left', va='top', color=textcolor)
    #SP
    xy = np.array([[165, 0], [360-70, 0], [360-70,-latmax]] )
    ax.plot(xy[:,0], xy[:,1], color=color, **kws)
    ax.text(360-70, -1, '70W', ha='right', va='top', color=textcolor)
    ##SA
    #xy = np.array([[360-70, 0], [360, 0], [360,-latmax]] )
    #ax.plot(xy[:,0], xy[:,1], color=color, **kws)

def wyplot_bars_cycle(basin='NA', ax=None, **kws):
    if ax is None:
        fig,ax = plt.subplots()
    color = kws.pop('color', 'C1') #TC color
    color_hu = 'C3' #HU color
    #basin = 'NA'
    pos = {
        'NA': [360-60, 52, 50, 20],
        'EP': [360-152, 40, 50, 20],
        'WP': [115, 55-2, 50, 20],
        'NI': [50, 40, 50, 20],
        'SI': [50, -60, 50, 20],
        'AU': [112, -60, 50, 20],
        'SP': [360-130, -50, 50, 20],
        'global': [360-50, -60, 50, 20],
        }
    axin = ax.inset_axes(pos[basin], transform=ax.transData)
    ds_cycle[basin].to_series().plot.bar(rot=0, ax=axin, color=color, **kws)
    ds_cycle_hu[basin].to_series().plot.bar(rot=0, ax=axin, color=color_hu, **kws)

    axin.grid(False)
    axin.set_xlabel('')
    axin.patch.set_alpha(0)
    axin.set_xticklabels(list('JFMAMJJASOND'))
    axin.set_ylim(0, ds_cycle[basin].max())
    yticks = np.arange(0, ds_cycle[basin].max(), 2) if ds_cycle[basin].max()>2 else np.arange(0, ds_cycle[basin].max(), 1)
    if basin == 'global': yticks = np.arange(0, ds_cycle[basin].max(), 4) #exception: higherinterval for global basin 
    axin.set_yticks(yticks)
    ntc = ds_cycle[basin].sum().item()
    axin.set_title(f'{basin} ({ntc:.1f})', color='C1')
    #axin color
    axcolor = 'gray'
    axin.spines['bottom'].set_color(axcolor)
    axin.spines['left'].set_color(axcolor)
    axin.tick_params(axis='x', colors=axcolor)
    axin.tick_params(axis='y', colors=axcolor)

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig,ax = plt.subplots(figsize=(11,6))
    ezplot = True if 'ezplot' in sys.argv else False
    alpha = 0.2 if ezplot else 1 #0.8
    ds.pipe(xtc.trackplot, ezplot=ezplot, alpha=alpha, lw=0.5)
    ds.isel(stage=0).plot.scatter('lon', 'lat', c='k', s=0.3,  zorder=10)

    ax.set_xticks(range(0, 361, 30))
    #mapplot(coastlines_width=1)
    mapplot(fill_continents=True, alpha=0.5)
    wyplot_basin_boundary(ax=ax)

    basin = 'NA'
    wyplot_bars_cycle(basin=basin, ax=ax)
    basin = 'EP'
    wyplot_bars_cycle(basin=basin, ax=ax)
    basin = 'WP'
    wyplot_bars_cycle(basin=basin, ax=ax)
    basin = 'NI'
    wyplot_bars_cycle(basin=basin, ax=ax)
    basin = 'SI'
    wyplot_bars_cycle(basin=basin, ax=ax)
    basin = 'AU'
    wyplot_bars_cycle(basin=basin, ax=ax)
    basin = 'SP'
    wyplot_bars_cycle(basin=basin, ax=ax)
    basin = 'global'
    wyplot_bars_cycle(basin=basin, ax=ax)

    ax.set_title(f'Observed TC genesis locations, tracks, intensities and frequency annual cycles over {yearspan.start}-{yearspan.stop}')
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_xlim(0,361)

    #edit colorbar
    if not ezplot:
        cbar = fig.axes[1]
        yticks = cbar.get_yticks()
        yy = (yticks[1:]+yticks[:-1])/2
        storm_types = ['TD', 'TS', 'CAT1', 'CAT2', 'CAT3', 'CAT4', 'CAT5']
        for y,storm in zip(yy, storm_types):
            cbar.text(110, y, storm, color='gray')
    
    #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()