#!/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
from IBTrACS.data_ntc_cycle import get_cycle as get_cycle_ibtracs
from ERA5.data_pxp_cycle import get_cycle as get_cycle_pxpERA5
from ERA5.data_ibntcxp_cycle import get_cycle as get_cycle_ibntcxpERA5
from ERA5.data_nseedxp_cycle_rainy import get_cycle as get_cycle_nseedxpERA5
from amipHadISST.data_ntc_cycle import get_cycle as get_cycle_hiram
from amipHadISST.data_pxp_cycle import get_cycle as get_cycle_pxp_hiram
from amipHadISST.data_ntcxp_cycle import get_cycle as get_cycle_ntcxp_hiram
from amipHadISST.data_nseedxp_cycle_rainy import get_cycle as get_cycle_nseedxp_hiram
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
basin = 'NA'
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, axes = plt.subplots(2, 1, figsize=(6, 6*9/16*1.5), sharex=True)
    alpha = 1/4
    marker = '.'
    ax = axes[0]
    labels1 = []
    # IBTrACS TC
    ds = get_cycle_ibtracs(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    label = 'IBTrACS TC'
    labels1.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)

    #ERA5 p x p
    ds = get_cycle_pxpERA5(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    label = 'ERA5 p$\\times$p'
    color = 'C1'
    ls = '--'
    labels1.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=color, alpha=alpha, zorder=10)
    ln1 = ds.mclim.plot(ax=ax, color=color, label=label, marker=marker, ls=ls)

    #IBTrACS ntc x ERA5 p 
    ds = get_cycle_ibntcxpERA5(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    label = 'IBTrACS N_TC$\\times$ ERA5 p'
    labels1.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)

    #ERA5 nseed x p
    ds = get_cycle_nseedxpERA5(basin=basin, years=years)
    ds = ds.sel(life=seed_olife)
    ds = ds/ds.mclim.sum('month')
    label = 'ERA5 N_SEED$\\times$p'
    labels1.append(label)
    color = 'C3'
    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)
    ln3 = ds.mclim.plot(ax=ax, color=color, label=label, marker=marker)

    ax = axes[1]
    labels2 = []
    # HiRAM TC
    ds = get_cycle_hiram(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    label = 'HiRAM 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)

    #HiRAM p x p
    ds = get_cycle_pxp_hiram(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    color = 'C1'
    ls = '--'
    label = 'HiRAM p$\\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=color, alpha=alpha, zorder=10)
    ln1 = ds.mclim.plot(ax=ax, color=color, label=label, marker=marker, ls=ls)

    #HiRAM ntcxp
    ds = get_cycle_ntcxp_hiram(basin=basin, years=years)
    #ds = ds.sel(life=seed_olife)
    ds = ds/ds.mclim.sum('month')
    label = 'HiRAM N_TC$\\times$p'
    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)

    #HiRAM nseed x p
    ds = get_cycle_nseedxp_hiram(basin=basin, years=years)
    ds = ds.sel(life=seed_olife)
    ds = ds/ds.mclim.sum('month')
    label = 'HiRAM N_SEED$\\times$p'
    labels2.append(label)
    color = 'C3'
    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)
    ln3 = ds.mclim.plot(ax=ax, color=color, label=label, marker=marker)

    ax = axes[0]
    ylabel = f'{basin} TC frequncy (normalized)'
    title_loc = 'left'
    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, 0.45)
    ax.set_ylabel(ylabel)
    ax.set_xticks(range(1,13))
    ax.set_xlabel('')
    ax.set_title('')
    ax.set_title('(a) Obs.', loc=title_loc)

    ax = axes[1]
    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, 0.45)
    ax.set_ylabel(ylabel)
    ax.set_xticks(range(1,13))
    ax.set_title('')
    ax.set_title('(b) HiRAM', loc=title_loc)

    if len(sys.argv)>1 and sys.argv[1]=='savefig':
        figname = __file__.replace('.py', '.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()