#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Mar  5 10:40:49 EST 2020
import sys, os.path, os, datetime
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
today = datetime.date.today().strftime('%Y-%m-%d')
print()

basin = 'NA'
years = slice('1981', '2018')
func_clim = lambda x: x.groupby('time.month').mean('time') # monthly climatology 
func_norm = lambda x: x/x.sum('month') # normalized by 12-month accumulation

fig, ax = plt.subplots(1, 1)
# ibtracs
ifile = '/tigress/wenchang/data/ibtracs/v04r00/analysis/ibtracs.counts.1980-2018.nc'
n_ib = xr.open_dataset(ifile)[basin]
n_ib.sel(time=years).pipe(func_clim).pipe(func_norm).plot(label='IBTrACS N$_{TC}$', color='k')

#HiRAM
ws = 17
ifile = f'data.ntc.{basin}.nc'
n_tc = xr.open_dataset(ifile)['N_TC']
n_tc.sel(time=years, ws=ws).mean('en').pipe(func_clim).pipe(func_norm).plot(label='HiRAM N$_{TC}$', color='gray')

#VI-predicted
VI0 = 0.02
ifile = f'data.p.{basin}.nc'
p = xr.open_dataset(ifile)['p']
p.sel(time=years, VI0=VI0).mean('en').pipe(func_clim).pipe(func_norm).plot(label='p(VI)', color='C0', ls='--')

#seed-predicted
life = 24 # hours
ifile = f'data.nseed.{basin}.nc'
n_seed = xr.open_dataset(ifile)['N_SEED']
n_seed.sel(time=years, life=life).mean('en').pipe(func_clim).pipe(func_norm).plot(label='N$_{seed}$', color='C0', ls=':')

#VI and seed combined
n_predicted = n_seed.sel(life=life, drop=True) * p.sel(VI0=VI0, drop=True).assign_coords(time=n_seed.time)
n_predicted.sel(time=years).mean('en').pipe(func_clim).pipe(func_norm).plot(label=r'N$_{seed}$$\times$p(VI)', color='C0', ls='-')

plt.legend(loc='upper left')
plt.ylabel(f'{basin} basin Normalized TC frequency')
plt.xticks(range(1,13))
#plt.grid('on')
plt.tight_layout()

figname = f'Fig.tcsp.norm.{basin}.png'
if os.path.exists(figname):
    figname = figname.replace('.png', f'.{today}.png')
plt.savefig(figname, dpi=120)

 
#if __name__ == '__main__':
#    tformat = '%Y-%m-%d %H:%M:%S'
#    t0 = datetime.datetime.now()
#    print('[start]:', t0.strftime(tformat))
#    
#    t1 = datetime.datetime.now()
#    print('[total time used]:', f'{(t1-t0).seconds:,} seconds')
#    print('[end]:', t1.strftime(tformat))
#    print()