#!/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
#
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_tracks = xr.open_dataset(ifile).sel(year=slice('1981', '2010'))
    stormHours = np.arange(6, 72+24+1, 6)
    wcHours = np.arange(0, 48+24+1, 6)
    wcWindContHours = np.arange(0, 48+1, 6) 
    windthd = 17
    ntc_glb = np.zeros((wcWindContHours.size, wcHours.size, stormHours.size)) + np.nan
    ntc_NA = np.zeros_like(ntc_glb)
    
    for iz in range(wcWindContHours.size):
        for iy in range(wcHours.size):
            if wcHours[iy] < wcWindContHours[iz]: continue
            #if iy==2: sys.exit()#test
            for ix in range(stormHours.size):
                if stormHours[ix] < wcHours[iy]: continue
                print(f'wcWindContHour = {wcWindContHours[iz]}; wcHour = {wcHours[iy]}; stormHour = {stormHours[ix]}')
                ds_filter = storm_filter(ds_tracks, stormHour=stormHours[ix], wcHour=wcHours[iy], wcWindContHour=wcWindContHours[iz], windthd=windthd)
                ntc_glb[iz, iy, ix] = tc_count(ds_filter).mean().item()*12
                ntc_NA[iz, iy, ix] = tc_count(ds_filter, basin='NA').mean().item()*12
                print(f'ntc_glb = {ntc_glb[iz, iy, ix]}; ntc_NA = {ntc_NA[iz, iy, ix]}')

    ntc_glb = xr.DataArray(ntc_glb, dims=['wcWindContHour', 'wcHour', 'stormHour'], coords=[wcWindContHours, wcHours, stormHours]) 
    ntc_NA = xr.DataArray(ntc_NA, dims=['wcWindContHour', 'wcHour', 'stormHour'], coords=[wcWindContHours, wcHours, stormHours]) 
    ds = xr.Dataset(dict(ntc_glb=ntc_glb, ntc_NA=ntc_NA))
    # save ds
    ds.to_netcdf(figdata)
    print(f'[saved]: {figdata}')

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, axes = plt.subplots(2, 1, figsize=(6, 6))
    wcWindContHour = 0
    ds_ = ds.sel(wcWindContHour=wcWindContHour)
    
    ax = axes[0]
    plt.sca(ax)
    da = ds_.ntc_glb
    da.plot(ax=ax)
    ax.set_title(f'{model} Global', loc='left')
    ax.set_xticks(da.stormHour.values)
    ax.set_yticks(da.wcHour.values)
    #ax.grid('on')
    ax.axhline(48, color='gray', ls='--')
    ax.axvline(72, color='gray', ls='--')
    for wcHour in da.wcHour.values:
        for stormHour in da.stormHour.values:
            n = da.sel(wcHour=wcHour, stormHour=stormHour).item()
            if n > 0:#not nan
                ax.text(stormHour, wcHour, f'{int(n)}', ha='center', va='center')
    
    ax = axes[1]
    plt.sca(ax)
    da = ds_.ntc_NA
    da.plot(ax=ax)
    ax.set_title(f'{model} NA', loc='left')
    ax.set_xticks(da.stormHour.values)
    ax.set_yticks(da.wcHour.values)
    #ax.grid('on')
    ax.axhline(48, color='gray', ls='--')
    ax.axvline(72, color='gray', ls='--')
    for wcHour in da.wcHour.values:
        for stormHour in da.stormHour.values:
            n = da.sel(wcHour=wcHour, stormHour=stormHour).item()
            if n > 0:#not nan
                ax.text(stormHour, wcHour, f'{int(n)}', ha='center', va='center')

    figname = __file__.replace('.py', f'_{tt.today()}.png')
    plt.savefig(figname)
    print(f'[saved]: {figname}')
    tt.check(f'**Done**')
    plt.show()