#!/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
from xtc.mask import get_landflag
import geoxarray
today = datetime.date.today().strftime('%Y-%m-%d')
print()

model = 'HIRAM'
expname = 'amipHadISST_tigercpu_intelmpi_18_540PE'
basin = 'NA'
seed_kinds = [6, 12, 18, 24] # least hours of life
idir = f'/tigress/wenchang/analysis/TC/{model}/{expname}/netcdf'
ifile = f'{idir}/{model}.{expname}.tc_tracks.allstorms.1971-2018.nc'
ds = (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))
    )
n_ocean = get_landflag(ds).pipe(lambda x: x<0.5).sum('stage') # number of over-ocean track points for each track
# calculate # of seed with different threshold of life hours
#n_seed = [ds.isel(stage=slice(seed_life//6, None)).pipe(tc_count, basin=basin)
n_seed = [ds.where(n_ocean>(seed_life//6)).pipe(tc_count, basin=basin)
    for seed_life in seed_kinds]
n_seed = xr.concat(n_seed, dim=pd.Index(seed_kinds, name='life'))
n_seed = n_seed.rename('N_SEED')

n_seed.mean('en').groupby('time.month').mean('time').plot(hue='life')
title = f'{model}.{expname}'
plt.title(title)
plt.tight_layout()
figname = f'fig.nseed.{basin}.png'
if os.path.exists(figname):
    figname = figname.replace('.png', f'.{today}.png')
plt.savefig(figname, dpi=120)

ds = n_seed.to_dataset()
ds.attrs['note'] = f'{model}.{expname}'

ofile = f'data.nseed.{basin}.nc'
if os.path.exists(ofile):
    print('[exists]:', ofile)
else:
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)
#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()