#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri May 29 11:37:36 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 misc.cim import cim, sem
# add /tigress/wenchang/analysis/seedTC to sys.path
thisdir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if sys.path[0] == '':
    if sys.path[1] != thisdir:
        sys.path.insert(1, thisdir)
else:
    if sys.path[0] != thisdir:
        sys.path.insert(0, thisdir)
print('sys.path[0:2]', sys.path[0:2])
from ERA5.data_nseedxp import get_nseedxp
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
def get_cycle(basin='NA', years=None):
    if years is None:
        years = slice('1980', '2018')
    da = get_nseedxp(basin=basin).sel(time=years)
    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
    basin = 'NA'
    years = slice('1980', '2018')
    label = os.path.basename(os.path.dirname(os.path.abspath(__file__))) # parent dir name
    figname = __file__.replace('.py', f'_{basin}_{label}_{tt.today()}.png')
    ds = get_cycle(basin=basin, years=years)
    fig, ax = plt.subplots()
    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'{label} N_SEED$\\times$p')
    plt.savefig(figname)
    print('[saved]:', figname)

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