#!/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
from misc.cim import cim
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

alpha = 0.3
fig, axes = plt.subplots(4, 1, sharex=True, figsize=(6,6))
ax = axes[0]
plt.sca(ax)
# ibtracs
ifile = '/tigress/wenchang/data/ibtracs/v04r00/analysis/ibtracs.counts.1980-2018.nc'
n_ib = xr.open_dataset(ifile)[basin].sel(time=years)
da = n_ib
m = da.groupby('time.month').mean()
e = da.groupby('time.month').apply(cim)
plt.fill_between(m.month, m-e, m+e, color='k', alpha=alpha)
m.plot(color='k', label='IBTrACS')
ax.set_title(''); ax.set_xlabel('')

#HiRAM
ws = 17
ifile = f'data.ntc.{basin}.nc'
n_tc = xr.open_dataset(ifile)['N_TC'].sel(time=years)
da = n_tc.sel(ws=17).stack(z=['en', 'time'])
m = da.groupby('time.month').mean()
e = da.groupby('time.month').apply(cim)
plt.fill_between(m.month, m-e, m+e, color='C0', alpha=alpha)
m.plot(color='C0', label='HiRAM')
ax.set_title(''); ax.set_xlabel('')
ax.set_ylabel('N$_{TC}$')
plt.legend(loc='upper left')

#p(VI)
ax = axes[1]
plt.sca(ax)
ifile = f'data.p.{basin}.nc'
p = xr.open_dataset(ifile)['p'].sel(time=years)
VI0 = 0.02
da = p.sel(VI0=VI0).stack(z=['en', 'time'])
m = da.groupby('time.month').mean()
e = da.groupby('time.month').apply(cim)
plt.fill_between(m.month, m-e, m+e, color='C0', alpha=alpha)
m.plot(color='C0', label=f'VI$_0$ = {VI0}')
VI0 = 0.04
da = p.sel(VI0=VI0).stack(z=['en', 'time'])
m = da.groupby('time.month').mean()
e = da.groupby('time.month').apply(cim)
plt.fill_between(m.month, m-e, m+e, color='C1', alpha=alpha)
m.plot(color='C1', label=f'VI$_0$ = {VI0}')
plt.legend(loc='upper left')
ax.set_title(''); ax.set_xlabel('')
ax.set_ylabel('p(VI)')

#seed
ax = axes[2]
plt.sca(ax)
ifile = f'data.nseed.{basin}.nc'
n_seed = xr.open_dataset(ifile)['N_SEED'].sel(time=years)
life = 24 # hours
da = n_seed.sel(life=life).stack(z=['en', 'time'])
m = da.groupby('time.month').mean()
e = da.groupby('time.month').apply(cim)
plt.fill_between(m.month, m-e, m+e, color='C0', alpha=alpha)
m.plot(color='C0', label=f'life>={life}h')
life = 12 # hours
da = n_seed.sel(life=life).stack(z=['en', 'time'])
m = da.groupby('time.month').mean()
e = da.groupby('time.month').apply(cim)
plt.fill_between(m.month, m-e, m+e, color='C1', alpha=alpha)
m.plot(color='C1', label=f'life>={life}h')
plt.legend(loc='upper left')
ax.set_title(''); ax.set_xlabel('')
ax.set_ylabel('N$_{seed}$')

# n_seed x p(VI)
ax = axes[3]
plt.sca(ax)
VI0, life = 0.02, 24
n_predicted = n_seed.sel(life=life, drop=True) * p.sel(VI0=VI0, drop=True).assign_coords(time=n_seed.time).sel(time=years)
da = n_predicted.stack(z=['en', 'time'])
m = da.groupby('time.month').mean()
e = da.groupby('time.month').apply(cim)
plt.fill_between(m.month, m-e, m+e, color='C0', alpha=alpha)
m.plot(color='C0', label=f'VI$_0$={VI0}, life={life}h')
VI0, life, color = 0.02, 12, 'C1'
n_predicted = n_seed.sel(life=life, drop=True) * p.sel(VI0=VI0, drop=True).assign_coords(time=n_seed.time).sel(time=years)
da = n_predicted.stack(z=['en', 'time'])
m = da.groupby('time.month').mean()
e = da.groupby('time.month').apply(cim)
plt.fill_between(m.month, m-e, m+e, color=color, alpha=alpha)
m.plot(color=color, label=f'VI$_0$={VI0}, life={life}h')
VI0, life, color = 0.04, 24, 'C2'
n_predicted = n_seed.sel(life=life, drop=True) * p.sel(VI0=VI0, drop=True).assign_coords(time=n_seed.time).sel(time=years)
da = n_predicted.stack(z=['en', 'time'])
m = da.groupby('time.month').mean()
e = da.groupby('time.month').apply(cim)
plt.fill_between(m.month, m-e, m+e, color=color, alpha=alpha)
m.plot(color=color, label=f'VI$_0$={VI0}, life={life}h')
VI0, life, color = 0.04, 12, 'C3'
n_predicted = n_seed.sel(life=life, drop=True) * p.sel(VI0=VI0, drop=True).assign_coords(time=n_seed.time).sel(time=years)
da = n_predicted.stack(z=['en', 'time'])
m = da.groupby('time.month').mean()
e = da.groupby('time.month').apply(cim)
plt.fill_between(m.month, m-e, m+e, color=color, alpha=alpha)
m.plot(color=color, label=f'VI$_0$={VI0}, life={life}h')
plt.legend(loc='upper left')
ax.set_title(''); #ax.set_xlabel('')
ax.set_ylabel(r'N$_{seed}\times$p(VI)')


plt.xticks(range(1,13))
#plt.grid('on')
plt.tight_layout()

#savefig
figname = f'fig.tcsp.spread.{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()