#!/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 = 'CTL1990s_v201910_tigercpu_intelmpi_18_540PE'
basin = 'NA'
# MDR lons and lats
lons = slice(360-80, 360-20)
lats = slice(10, 20)
seed_life = [0, 12, 18, 24][2] #hours
figname = f'fig.{model}.{expname.split("_")[0]}.spTS.{basin}.{seed_life}h.png'
dpi = 100
idir = f'/tigress/wenchang/analysis/TC/{model}/{expname}/netcdf'
ifile = f'{idir}/{model}.{expname}.tc_tracks.TS.0101-0150.nc'
n_tc = xr.open_dataset(ifile).pipe(lambda ds: ds.where(ds.windmax.max('stage')>17)).pipe(tc_count, basin)
ifile = f'{idir}/{model}.{expname}.tc_tracks.allstorms.0101-0150.nc'
n_seed = (xr.open_dataset(ifile) 
    # genesis locations within the tropics
    .pipe(lambda x: x.where(x.lat.isel(stage=0)>-30).where(x.lat.isel(stage=0)<30))
    # select tracks of at least 24h long
    .isel(stage=slice(seed_life//6, None))
    .pipe(tc_count, basin)
    )
ifiles = [f'/tigress/wenchang/MODEL_OUT/{model}/{expname}/analysis_wy/TCI/{year:04d}0101.atmos_month.VI.nc'
    for year in range(101, 150+1)
    ]
# ventilation index
vi = (xr.open_mfdataset(ifiles)['VI']
    .rename(dict(grid_xt='lon', grid_yt='lat'))
     # MDR region selection
    .sel(lat=lats, lon=lons)
    .geo.fldmean().load() 
    .assign_coords(time=n_tc.time)
    )

ds = xr.Dataset(dict(
    n_TC=n_tc, 
    n_seed=n_seed, 
    VI=vi, 
    p=1/(1+vi/0.02), 
    n_seed_x_p=n_seed/(1+vi/0.02)
    ))
ds_clim = ds.groupby('time.month').mean('time')
ds_clim['n_seedClim_x_pClim'] = ds_clim['n_seed']/(1+ds_clim['VI']/0.02)

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} TS, seed and VI, seed_life={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', 'n_seedClim_x_pClim']].to_dataframe().plot(secondary_y=['n_seed_x_p', 'n_seedClim_x_pClim'], 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()