#!/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
# add /tigress/wenchang/analysis/seedTC to sys.path
maindir = '/tigress/wenchang/analysis/seedTC'
if maindir not in sys.path:
    sys.path.append(maindir)
from AM2p5.AU.data_nseed import get_nseed
from AM2p5.AU.data_pTang import get_pTang
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
def get_nseedxp(basin='NA'):
    '''get n_seed \times p given n_seed and p'''
    unify_time = lambda x: x.assign_coords(time=x.indexes['time'].to_datetimeindex()).resample(time='MS').mean()
    nseed = get_nseed(basin=basin).pipe(unify_time)
    pTang = get_pTang(basin=basin).pipe(unify_time)
    
    return nseed*pTang
    
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    basin = 'AU'
    #label = os.path.basename(os.path.dirname(os.path.abspath(__file__))) # parent dir name
    source = 'AM2.5'
    label = f'{source} {basin}'
    figname = __file__.replace('.py', f'_{basin}_{source}_{tt.today()}.png')
    da = get_nseedxp(basin=basin)
    fig, ax = plt.subplots()
    da.groupby('time.year').sum('time').mean('en').plot(hue='life', ax=ax)
    ax.set_ylabel(f'{label} N_SEED $\\times$ p')

    plt.savefig(figname)
    print('[saved]:', figname)

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