#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu May 21 11:24:22 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.mask import get_landflag
from misc.cim import cim, sem #confidence interval of the mean
#
if __name__ == '__main__':
    tt.check('end import')
    
#start from here
thisdir = os.path.dirname(__file__)
source = 'AM2.5C360_amipHadISST'
basin = 'NA'
def get_nseed(basin='NA'):
    #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/rainyseed/AM2.5C360_amipHadISST_tigercpu_intelmpi_18_1080PE_en01-03_yr1971to2018_mo1to12_rainPerc99_distThresh2.9_sizeThresh5.1_latLim30.nc'
    #ofile = os.path.join(thisdir, f'data_nseed_rainy2xvort_{basin}.nc')
    ofile = __file__.replace('.py', f'_{basin}.nc')
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        return xr.open_dataset(ofile)['N_SEED']
    seed_kinds = [6, 12, 18, 24] # least hours of life
    """
    ds = (xr.open_dataset(ifile)
        # genesis locations within the tropics
        .pipe(lambda x: x.where(x.lat.isel(stage=0)>-30).where(x.lat.isel(stage=0)<30))
    )
    """
    ds = xr.open_dataset(ifile).rename(track='storm', lifetime='stage')
    #ds = ds.where(ds.vort.max('stage')>8e-4)
    ds = ds.where(ds.vort.max('stage')>12e-4)
    n_ocean = get_landflag(ds).pipe(lambda x: x<0.5).sum('stage')
    # calculate # of seed with different threshold of life hours
    n_seed = [ds.where(n_ocean>(seed_life//6)).pipe(tc_count, basin=basin)
        for seed_life in seed_kinds]
    #if 'vmax' in ds and 'windmax' not in ds:
    #    ds = ds.rename(vmax='windmax')
    n_seed = xr.concat(n_seed, dim=pd.Index(seed_kinds, name='life'))
    n_seed = n_seed.rename('N_SEED')

    #save to nc
    n_seed.to_dataset().to_netcdf(ofile)
    print('[saved]:', ofile)

    return n_seed

#annual cycle
years = slice('1980', '2018')
def get_cycle(basin='NA', years=years):
    da = get_nseed(basin=basin)
    da = da.sel(time=years)
    # stack dims of ['en', 'time'] and use the stacked 'time' as the new coord
    if 'en' in da.dims:
        da = da.stack(s=['en', 'time']).pipe(lambda x: x.assign_coords(s=x.time)).rename(s='time')
    da_mclim = da.groupby('time.month').mean('time')
    da_cim = da.groupby('time.month').map(cim, dim='time')
    da_sem = da.groupby('time.month').map(sem, dim='time')
    ds_mclim = xr.Dataset(dict(
        mclim=da_mclim,
        cim=da_cim,
        sem=da_sem
        ))
    return ds_mclim

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, ax = plt.subplots()
    #figname = f'data_nseed_rainy2xvort_{basin}_{tt.today()}.png'
    figname = __file__.replace('.py', f'_{basin}_{tt.today()}.png')
    da = get_nseed(basin=basin)
    da.groupby('time.year').sum('time').mean('en').plot(hue='life')
    plt.title(f'{basin} N_SEED(rainy): {source}')
    """
    figname = f'data_nseed_cycle_{basin}_{source}_{tt.today()}.png'
    ds_ = get_cycle(basin=basin, years=years)
    long_name = 'N_SEED(NA, {source})'
    for i,life in enumerate(ds_.life):
        ax.bar(x='month', height='mclim', width=0.6, yerr='cim', data=ds_.sel(life=life).assign_coords(month=ds_.month.values+0.1*i), label=f'{life.item()}')    
    ax.legend()
    ax.set_xlabel('month')
    ax.set_ylabel(f'{source} seed # climatology')
    ax.set_xticks(range(1,13))
    """
    plt.savefig(figname)
    print('[saved]:', figname)

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