#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Oct 16 16:36:44 EDT 2020
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
from xtc import tc_count
#from xtc.filter import storm_filter
from xtc.filter import selTS
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
model = 'HiRAM'
figdata = __file__.replace('.py', '.nc')
if os.path.exists(figdata):
    ds = xr.open_dataset(figdata)
    print(f'[loaded]: {figdata}')
else:
    ifile = '/tigress/wenchang/analysis/TC/HIRAM/amipHadISST_tigercpu_intelmpi_18_540PE/netcdf/HIRAM.amipHadISST_tigercpu_intelmpi_18_540PE.tc_tracks.allstorms.1971-2018.nc'
    ds = xr.open_dataset(ifile).sel(year=slice('1981', '2010'))
    stormHours = np.arange(72-36, 72+24+1, 6)
    wcHours = np.arange(48-36, 48+24+1, 6)
    wcWindContHours = np.arange(36-36, 36+24+1, 6)
    ntc15_glb = np.zeros_like(stormHours)
    ntc17_glb = np.zeros_like(stormHours)
    ntc15_NA = np.zeros_like(stormHours)
    ntc17_NA = np.zeros_like(stormHours)
    
    for i in range(stormHours.size):
        print(f'{i+1:02d} of {stormHours.size:02d}: stormHour = {stormHours[i]}')
        #ds15 = storm_filter(ds, stormHour=stormHours[i], wcHour=wcHours[i], wcWindContHour=wcWindContHours[i], windthd=15.75)
        #ds17 = storm_filter(ds, stormHour=stormHours[i], wcHour=wcHours[i], wcWindContHour=wcWindContHours[i], windthd=17)
        ds15 = selTS(ds, stormHour=stormHours[i], wcHour=wcHours[i], wcWindContHour=wcWindContHours[i], windthd=15.75)
        ds17 = selTS(ds, stormHour=stormHours[i], wcHour=wcHours[i], wcWindContHour=wcWindContHours[i], windthd=17)
        ntc15_glb[i] = tc_count(ds15, ws=17).mean().item()*12
        ntc17_glb[i] = tc_count(ds17, ws=17).mean().item()*12
        ntc15_NA[i] = tc_count(ds15, basin='NA', ws=17).mean().item()*12
        ntc17_NA[i] = tc_count(ds17, basin='NA', ws=17).mean().item()*12

    ntc15_glb = xr.DataArray(ntc15_glb, dims='stormHour', coords=[stormHours]) 
    ntc17_glb = xr.DataArray(ntc17_glb, dims='stormHour', coords=[stormHours]) 
    ntc15_NA = xr.DataArray(ntc15_NA, dims='stormHour', coords=[stormHours]) 
    ntc17_NA = xr.DataArray(ntc17_NA, dims='stormHour', coords=[stormHours]) 
    ds = xr.Dataset(dict(ntc15_glb=ntc15_glb, ntc17_glb=ntc17_glb, ntc15_NA=ntc15_NA, ntc17_NA=ntc17_NA))
    # save ds
    ds.to_netcdf(figdata)
    print(f'[saved]: {figdata}')

stormHours = ds.stormHour
ntc15_glb = ds.ntc15_glb
ntc17_glb = ds.ntc17_glb
ntc15_NA = ds.ntc15_NA
ntc17_NA = ds.ntc17_NA

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, axes = plt.subplots(2, 1, figsize=(6, 6))
    
    ax = axes[0]
    plt.sca(ax)
    ax.plot(stormHours, ntc15_glb, 'o-', label='WC wind threshold: 15.75 m/s')
    ax.plot(stormHours, ntc17_glb, '^-', label='WC wind threshold: 17 m/s')
    ax.legend()
    ax.set_title(f'{model} Global', loc='left')
    ax.axhline(86, ls='--', color='gray')
    ax.set_xticks(stormHours)
    ax.set_ylabel('N_TC')
    ax.text(0.99, 0.5, 'WC_life = 48 $+$ (storm_life $-$ 72)\nWC_wind_cont = 36 $+$ (storm_life $-$ 72)',
        transform=ax.transAxes, ha='right', va='center')
    
    ax = axes[1]
    plt.sca(ax)
    ax.plot(stormHours, ntc15_NA, 'o-', label='WC wind threshold: 15.75 m/s')
    ax.plot(stormHours, ntc17_NA, '^-', label='WC wind threshold: 17 m/s')
    ax.legend()
    ax.set_title(f'{model} NA', loc='left')
    ax.axhline(12, ls='--', color='gray')
    ax.set_xticks(stormHours)
    ax.set_xlabel('storm life [hours]')
    ax.set_ylabel('N_TC')

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        #figname = __file__.replace('.py', f'_{tt.today()}.png')
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
        #print(f'[saved]: {figname}')

    tt.check(f'**Done**')
    plt.show()