#!/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
maindir = '/tigress/wenchang/analysis/seedTC'
if maindir not in sys.path:
    sys.path.append(maindir)
from amipHadISST.data_ntc import get_ntc
from misc.cim import cim, sem #confidence interval of the mean
#
if __name__ == '__main__':
    tt.check('end import')
    
#start from here
#yearly 
def get_yearly(basin='NA', months=None, normalize=True):
    years_base = range(1981,2011)
    if months is None:
        if basin == 'NA':
            months = [8, 9, 10]
    da = get_ntc(basin=basin)
    da_yearly = da.where(da.time.dt.month.isin(months)).groupby('time.year').sum('time')
    if normalize:
        da_base = da_yearly.sel(year=years_base).mean('year')
        if 'en' in da_base.dims:
            da_base = da_base.mean('en')
        da_yearly = (da_yearly - da_base)/da_base
            
    return da_yearly

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    basin = 'NA'
    source = 'HiRAM_amipHadISST'
    figname = __file__.replace('.py', f'_{basin}_{tt.today()}.png')
    da = get_yearly(basin=basin)
    long_name = f'N_TC fraction({basin}, {source})'
    if 'en' in da.dims:
        da.plot(hue='en', color='C0', alpha=0.2, lw=1, add_legend=False)
        da.mean('en').plot(marker='o', fillstyle='none')
    else:
        da.plot(marker='o', fillstyle='none')
    plt.title(long_name)
    plt.savefig(figname)
    print('[saved]:', figname)

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