#!/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_yearly = lambda x: x.where((x.time.dt.month==8)|(x.time.dt.month==9)).groupby('time.year').sum('time') # monthly climatology 
#func_norm = lambda x: (x-x.mean('year'))/x.std('year') # normalized by 12-month accumulation
func_norm = lambda x: (x - x.mean('year'))/x.mean('year') # normalized by multi-year mean 

fig, ax = plt.subplots(1, 1, figsize=(8, 4))
# 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_yearly).pipe(func_norm).plot(label='IBTrACS', 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_yearly).pipe(func_norm).plot(label='model', 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_yearly).pipe(func_norm).plot(label='VI predicted', color='C0')

#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_yearly).pipe(func_norm).plot(label='seed predicted', color='C1')

#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_yearly).pipe(func_norm).plot(label='seedxp(VI) predicted', color='C2')

plt.legend()
plt.ylabel('Normalized TC frequency')
plt.tight_layout()

figname = f'Fig.tcsp.yearly.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()