#!/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
import xaddon
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(1, 3, sharey=True, figsize=(8,3))
ax = axes[0]
plt.sca(ax)

#HiRAM ntc vs. p(VI)
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)
y = da.groupby('time.month').mean('time').mean('en')
#ax.set_title(''); ax.set_xlabel('')
#ax.set_ylabel('N$_{TC}$')
#p(VI)
ifile = f'data.p2.{basin}.nc'
p = xr.open_dataset(ifile)['p'].sel(time=years)
VI0 = 0.02
da = p.sel(VI0=VI0)
x = da.groupby('time.month').mean('time').mean('en')
#ax.set_title(''); ax.set_xlabel('')
#ax.set_ylabel('p(VI)')
#plt.scatter(x, y)
df = pd.DataFrame(dict(x=x, y=y, month=x.month))
df.plot.scatter(x='x', y='y', ax=ax)#, c='month', cmap='twilight', ax=ax)
rg = y.linregress.on(x)
ax.plot(x, rg.predicted)
ax.text(0.02, 1-0.02, f'y={rg.slope.item():.1f}*x $-$ {-rg.intercept.item():.2f}\nr$^2$={rg.r.item()**2:.2f}', 
    ha='left', va='top', transform=ax.transAxes)
ax.set_xlim(0, None)
ax.set_xlabel('p(VI): 1/(1+(VI/VI$_0$)$^2$)')
ax.set_ylabel(f'{basin} N$_{{TC}}$')

#seed
ax = axes[1]
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)
x = da.groupby('time.month').mean('time').mean('en')
df = pd.DataFrame(dict(x=x, y=y, month=x.month))
df.plot.scatter(x='x', y='y', ax=ax)#, c='month', cmap='twilight', ax=ax)
rg = y.linregress.on(x)
ax.plot(x, rg.predicted)
ax.text(0.02, 1-0.02, f'y={rg.slope.item():.1f}*x $-$ {-rg.intercept.item():.2f}\nr$^2$={rg.r.item()**2:.2f}', 
    ha='left', va='top', transform=ax.transAxes)
ax.set_xlim(0, None)
ax.set_xlabel('N$_{seed}$')

# n_seed x p(VI)
ax = axes[2]
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
x = da.groupby('time.month').mean('time').mean('en')
df = pd.DataFrame(dict(x=x, y=y, month=x.month))
df.plot.scatter(x='x', y='y', ax=ax)#, c='month', cmap='twilight', ax=ax)
rg = y.linregress.on(x)
ax.plot(x, rg.predicted)
ax.text(0.02, 1-0.02, f'y={rg.slope.item():.1f}*x $-$ {-rg.intercept.item():.2f}\nr$^2$={rg.r.item()**2:.2f}', 
    ha='left', va='top', transform=ax.transAxes)
ax.set_xlim(0, None)
ax.set_ylim(0, None)
ax.set_xlabel(r'N$_{seed}\times$p(VI)')

#plt.grid('on')
plt.tight_layout()

#savefig
figname = f'Fig.scatter.{basin}.p2.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()