#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri May 29 16:26:38 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
#maindir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
maindir = '/tigress/wenchang/analysis/seedTC'
if maindir not in sys.path:
    sys.path.append(maindir)
from AM2p5.SP.data_ntc_cycle import get_cycle as get_cycle
from AM2p5.SP.data_pTang_cycle import get_cycle as get_cycle_p
from AM2p5.SP.data_nseed_cycle_rainy import get_cycle as get_cycle_nseed
from AM2p5.SP.data_nseedxp_cycle_rainy import get_cycle as get_cycle_nseedxp
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
basin = 'SP'
years = slice('1980', '2018')
seed_olife = 12
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    figname = __file__.replace('.py', f'_{basin}_{tt.today()}.png')
    #fig, ax = plt.subplots()
    fig, ax = plt.subplots()
    marker = '.'
    alpha = 1/4
    labels2 = []
    # AM2.5 TC
    ds = get_cycle(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    label = 'AM2.5 TC'
    labels2.append(label)
    b1 = ax.bar(x='month', height='mclim', data=ds, label=label)
    ax.errorbar(x='month', y='mclim', yerr='cim', capsize=2, ls='none', color='k', data=ds, label=None)

    #AM2.5 p
    ds = get_cycle_p(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    label = 'AM2.5 p'
    labels2.append(label)
    color = 'C1'
    ls = '--'
    ds['lower'] = ds.mclim - ds.cim
    ds['upper'] = ds.mclim + ds.cim
    ax.fill_between(x='month', y1='lower', y2='upper', data=ds, color=color, alpha=alpha, zorder=10)
    ln1 = ds.mclim.plot(ax=ax, color=color, label=label, marker=marker, ls=ls)

    #AM2.5 nseed
    ds = get_cycle_nseed(basin=basin, years=years)
    ds = ds.sel(life=seed_olife)
    ds = ds/ds.mclim.sum('month')
    label = 'AM2.5 N_SEED(rainy)'
    labels2.append(label)
    color = 'C2'
    ls = ':'
    ds['lower'] = ds.mclim - ds.cim
    ds['upper'] = ds.mclim + ds.cim
    ax.fill_between(x='month', y1='lower', y2='upper', data=ds, color=color, alpha=alpha, zorder=10)
    ln2 = ds.mclim.plot(ax=ax, color=color, label=label, marker=marker, ls=ls)

    #AM2.5 nseed x p
    ds = get_cycle_nseedxp(basin=basin, years=years)
    ds = ds.sel(life=seed_olife)
    ds = ds/ds.mclim.sum('month')
    label = 'AM2.5 N_SEED(rainy)$\\times$p'
    labels2.append(label)
    ds['lower'] = ds.mclim - ds.cim
    ds['upper'] = ds.mclim + ds.cim
    ax.fill_between(x='month', y1='lower', y2='upper', data=ds, color='C3', alpha=alpha, zorder=10)
    ln3 = ds.mclim.plot(ax=ax, color='C3', label=label, marker=marker)

    ylabel = f'{basin} TC frequncy (normalized)'  
    legend_raw = ax.legend()
    handles = legend_raw.legendHandles
    labels = [t.get_text() for t in legend_raw.get_texts()]
    ax.legend(handles[-1:-2:-1]+handles[:-1], labels[-1:-2:-1]+labels[:-1])#move the bar legend from bottom to top
    ax.set_ylim(0, None)
    ax.set_ylabel(ylabel)
    ax.set_xticks(range(1,13))
    ax.set_title('')
    #ax.set_title('HiRAM', loc='left')

    plt.savefig(figname)
    print('[saved]:', figname)
    tt.check(f'**Done**')
    plt.show()