#!/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.density.TS17.1980-2021.yearly.nc'
yearspan = slice(1980,2021)
ds = xr.open_dataset(ifile).sel(year=yearspan).load()
 
#TC
ifile = 'IBTrACS.ALL.v04r00.2022-01-25.counts.TS17.1980-2021.yearly.nc'
#ifile = 'IBTrACS.ALL.v04r00.2022-01-25.counts.TS17.2days17.1980-2021.yearly.nc'
ds_ntc = xr.open_dataset(ifile).assign_coords(month=range(1,12+1))
#HU
ifile = 'IBTrACS.ALL.v04r00.2022-01-25.counts.TS33.1980-2021.yearly.nc'
ds_nhu = 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_lines_yearly(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_ntc[basin].plot(ax=axin, color=color, **kws)
    ds_nhu[basin].plot(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_ntc[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_ntc[basin].mean().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)
    axin.set_ylabel('')

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig,ax = plt.subplots(figsize=(11,6))
    levels = np.logspace(-2,6,9,base=2)
    ds['density'].mean('year', keep_attrs=True).plot.contourf(levels=levels, cmap='Purples')

    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_lines_yearly(basin=basin, ax=ax)
    basin = 'EP'
    wyplot_lines_yearly(basin=basin, ax=ax)
    basin = 'WP'
    wyplot_lines_yearly(basin=basin, ax=ax)
    basin = 'NI'
    wyplot_lines_yearly(basin=basin, ax=ax)
    basin = 'SI'
    wyplot_lines_yearly(basin=basin, ax=ax)
    basin = 'AU'
    wyplot_lines_yearly(basin=basin, ax=ax)
    basin = 'SP'
    wyplot_lines_yearly(basin=basin, ax=ax)
    basin = 'global'
    wyplot_lines_yearly(basin=basin, ax=ax)

    ax.set_title(f'Observed TC track density climatology and frequency interannual variability')
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_xlim(0,361)
    ax.set_ylim(-80,80)

    #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()