#!/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')

alpha = 0.3
fig, axes = plt.subplots(1, 2, sharey=True, figsize=(8,3))
# regression
ax = axes[0]
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)
y = da.groupby('time.month').mean('time').mean('en')
#y = np.log(y)
ifile = f'data.omgfstar.{basin}.nc'
da = xr.open_dataarray(ifile).sel(time=years)
x = da.groupby('time.month').mean('time').mean('en')
x = x/1e7
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_ylabel('log(N$_{seed}$)')
ax.set_ylabel('N$_{seed}$')
ax.set_xlabel(da.attrs['long_name'] + ' [1e7 ' + da.attrs['units'] + ']')

# annual cyclej
ax = axes[1]
plt.sca(ax)
y.plot(ax=ax, label='counted')
y_ = (da/1e7*rg.slope + rg.intercept).groupby('time.month').mean('time').mean('en')
y_.plot(ax=ax, label='predicted')
ax.legend()
ax.set_title('')

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

#savefig
figname = f'Fig.scatter.nseed_vs_omgfstar.{basin}.png'
if os.path.exists(figname):
    figname = figname.replace('.png', f'.{today}.png')
plt.savefig(figname)

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