#!/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'
    #ifile = '/tigress/wenchang/analysis/TC/AM2.5C360/amipHadISST_tigercpu_intelmpi_18_1080PE/netcdf/tc_tracks.allstorms.1971-2018.nc'
    #ifile = '/tigress/wenchang/analysis/TC/AM2.5/amipHadISST/netcdf/tc_tracks.allstorms.1971-2018.nc'
    ds = xr.open_dataset(ifile).sel(year=slice('1981', '2010'))
    minWindmax_vec = np.arange(17,34)
    dHour = 12 #hour offset from the standard
    #stormHours = np.arange(72-36, 72+24+1, 6)
    stormHour = 72 + dHour
    #wcHours = np.arange(48-36, 48+24+1, 6)
    wcHour = 48 + dHour
    #wcWindContHours = np.arange(36-36, 36+24+1, 6)
    wcWindContHour = 36 + dHour

    ntc15_glb = np.zeros_like(minWindmax_vec)
    ntc17_glb = np.zeros_like(minWindmax_vec)
    ntc15_NA = np.zeros_like(minWindmax_vec)
    ntc17_NA = np.zeros_like(minWindmax_vec)
    
    ds15 = selTS(ds, stormHour=stormHour, wcHour=wcHour, wcWindContHour=wcWindContHour, windthd=15.75)
    ds17 = selTS(ds, stormHour=stormHour, wcHour=wcHour, wcWindContHour=wcWindContHour, windthd=17)
    for i in range(minWindmax_vec.size):
        print(f'{i+1:02d} of {minWindmax_vec.size:02d}: minWndmax = {minWindmax_vec[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)
        ntc15_glb[i] = tc_count(ds15, ws=minWindmax_vec[i]).mean().item()*12
        ntc17_glb[i] = tc_count(ds17, ws=minWindmax_vec[i]).mean().item()*12
        ntc15_NA[i] = tc_count(ds15, basin='NA', ws=minWindmax_vec[i]).mean().item()*12
        ntc17_NA[i] = tc_count(ds17, basin='NA', ws=minWindmax_vec[i]).mean().item()*12

    ntc15_glb = xr.DataArray(ntc15_glb, dims='minWindmax', coords=[minWindmax_vec]) 
    ntc17_glb = xr.DataArray(ntc17_glb, dims='minWindmax', coords=[minWindmax_vec]) 
    ntc15_NA = xr.DataArray(ntc15_NA, dims='minWindmax', coords=[minWindmax_vec]) 
    ntc17_NA = xr.DataArray(ntc17_NA, dims='minWindmax', coords=[minWindmax_vec]) 
    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}')

minWindmax = ds.minWindmax
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(minWindmax, ntc15_glb, 'o-', label='WC wind threshold: 15.75 m/s')
    ax.plot(minWindmax, ntc17_glb, '^-', label='WC wind threshold: 17 m/s')
    ax.legend()
    ax.set_title(f'{model.replace("2p5", "2.5")} Global', loc='left')
    ax.axhline(47, ls='--', color='gray')
    ax.set_xticks(minWindmax)
    ax.set_ylabel('N_HU')
    
    ax = axes[1]
    plt.sca(ax)
    ax.plot(minWindmax, ntc15_NA, 'o-', label='WC wind threshold: 15.75 m/s')
    ax.plot(minWindmax, ntc17_NA, '^-', label='WC wind threshold: 17 m/s')
    ax.legend()
    ax.set_title(f'{model.replace("2p5", "2.5")} NA', loc='left')
    ax.axhline(6, ls='--', color='gray')
    ax.set_xticks(minWindmax)
    ax.set_xlabel('minWindmax [m/s]')
    ax.set_ylabel('N_HU')

    #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()