#!/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 selTS
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
model = 'HiRAM'
figdata = __file__.replace('.py', '.nc')
stormHour = 72 
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'
    #ifile = '/tigress/wenchang/analysis/TC/AM4_urban/amip1870_urban_luh2_tigercpu_intelmpi_18_576PE/netcdf/tc_tracks.allstorms.1870-2014.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(wcHours)
    ntc17_glb = np.zeros_like(wcHours)
    ntc15_NA = np.zeros_like(wcHours)
    ntc17_NA = np.zeros_like(wcHours)
    
    for i in range(wcHours.size):
        print(f'{i+1:02d} of {wcHours.size:02d}: wcHour = {wcHours[i]}')
        ds15 = selTS(ds, stormHour=stormHour, wcHour=wcHours[i], wcWindContHour=wcWindContHours[i], windthd=15.75)
        ds17 = selTS(ds, stormHour=stormHour, 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='wcHour', coords=[wcHours]) 
    ntc17_glb = xr.DataArray(ntc17_glb, dims='wcHour', coords=[wcHours]) 
    ntc15_NA = xr.DataArray(ntc15_NA, dims='wcHour', coords=[wcHours]) 
    ntc17_NA = xr.DataArray(ntc17_NA, dims='wcHour', coords=[wcHours]) 
    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}')

wcHours = ds.wcHour
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(wcHours, ntc15_glb, 'o-', label='WC wind threshold: 15.75 m/s')
    ax.plot(wcHours, 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(86, ls='--', color='gray')
    ax.set_xticks(wcHours)
    ax.set_ylabel('N_TC')
    ax.text(0.99, 0.5, f'total_life_threshold =  {stormHour}\nWC_wind_cont = wcHour $-$ 12',
        transform=ax.transAxes, ha='right', va='center')
    
    ax = axes[1]
    plt.sca(ax)
    ax.plot(wcHours, ntc15_NA, 'o-', label='WC wind threshold: 15.75 m/s')
    ax.plot(wcHours, 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(12, ls='--', color='gray')
    ax.set_xticks(wcHours)
    ax.set_xlabel('warm core accumulative threshold [hours]')
    ax.set_ylabel('N_TC')

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