#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Feb 20 15:52:22 EST 2020
import sys, os.path, os, datetime
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
from xtc import tc_count
import geoxarray
print()

model = 'HIRAM'
expname = 'amipHadISST_tigercpu_intelmpi_18_540PE'
basin = 'NA'
seed_life = 18 
VI0 = 0.02
ws = 17 
years = slice('1986', '2005')
ifile = f'data.ntc.{basin}.nc'
n_tc = xr.open_dataset(ifile)['N_TC'].sel(ws=ws, drop=True)

ifile = f'data.nseed.{basin}.nc'
n_seed = xr.open_dataset(ifile)['N_SEED'].sel(life=seed_life, drop=True)

ifile = f'data.p.{basin}.nc'
p = xr.open_dataset(ifile)['p'].assign_coords(time=n_tc.time).sel(VI0=VI0, drop=True)

ds = xr.Dataset(dict(
    n_TC=n_tc, 
    n_seed=n_seed, 
    p=p,
    n_seed_x_p=n_seed*p,
    ))
ds_clim = ds.sel(time=years).groupby('time.month').mean(['time', 'en'])

fig, axes = plt.subplots(3, 1, figsize=(6, 6), sharex=True)
ax = axes[0]
ds_clim[['n_TC', 'n_seed']].to_dataframe().plot(secondary_y='n_seed', ax=ax)
ax.set_title(f'{model} {basin} TC({ws}), p({VI0}) and seed({seed_life}h)')
 
ax = axes[1]
ds_clim[['n_TC', 'p']].to_dataframe().plot(secondary_y='p', ax=ax)

ax = axes[2]
ds_clim[['n_TC', 'n_seed_x_p' ]].to_dataframe().plot(secondary_y=['n_seed_x_p',], ax=ax)

plt.tight_layout()
#plt.savefig(figname, dpi=dpi)
#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()