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

model = 'HIRAM'
expname = 'amipHadISST_tigercpu_intelmpi_18_540PE'
basin = 'NA'
idir = f'/tigress/wenchang/analysis/TC/{model}/{expname}/netcdf'
ifile = f'{idir}/{model}.{expname}.tc_tracks.TS.1971-2018.nc'
wss = [17, 33]
n_tc = [xr.open_dataset(ifile).pipe(tc_count, basin=basin, ws=ws) for ws in wss]
n_tc = xr.concat(n_tc, dim=pd.Index(wss, name='ws'))
n_tc.name = 'N_TC'

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

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

ofile = f'data.ntc.{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()