#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri May 29 16:26:38 EDT 2020
# v2 (2020-10-13): add AM2.5 and AM2.5C360
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_pTang_cycle import get_cycle as get_cycle_pERA5
from ERA5.data_nseed_cycle_rainy import get_cycle as get_cycle_nseedERA5
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_pTang_cycle import get_cycle as get_cycle_p_hiram
from amipHadISST.data_nseed_cycle_rainy import get_cycle as get_cycle_nseed_hiram
from amipHadISST.data_nseedxp_cycle_rainy import get_cycle as get_cycle_nseedxp_hiram
from AM2p5.data_ntc_cycle import get_cycle as get_cycle_floram
from AM2p5.data_pTang_cycle import get_cycle as get_cycle_p_floram
from AM2p5.data_nseed_cycle_rainy import get_cycle as get_cycle_nseed_floram
from AM2p5.data_nseedxp_cycle_rainy import get_cycle as get_cycle_nseedxp_floram
from AM2p5C360.data_ntc_cycle import get_cycle as get_cycle_floramplus
from AM2p5C360.data_pTang_cycle import get_cycle as get_cycle_p_floramplus
from AM2p5C360.data_nseed_cycle_rainy import get_cycle as get_cycle_nseed_floramplus
from AM2p5C360.data_nseedxp_cycle_rainy import get_cycle as get_cycle_nseedxp_floramplus
#
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, 2, figsize=(8, 6*9/16*1.5), sharex=False, sharey=True)
    marker = '.'
    alpha = 1/4

    ax = axes[0, 0]
    labels1 = []
    # IBTrACS TC
    ds = get_cycle_ibtracs(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    label = 'IBTrACS N_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
    ds = get_cycle_pERA5(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    label = 'ERA5 p($\Lambda$)'
    labels1.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)
    #ln1 = ax.plot('month', 'mclim', color='C1', label=label, data=ds)

    #ERA5 nseed
    ds = get_cycle_nseedERA5(basin=basin, years=years)
    ds = ds.sel(life=seed_olife)
    ds = ds/ds.mclim.sum('month')
    label = 'ERA5 N_SEED'
    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($\Lambda$)'
    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[0, 1]
    labels2 = []
    # HiRAM TC
    ds = get_cycle_hiram(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    label = 'N_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
    ds = get_cycle_p_hiram(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    label = 'p($\Lambda$)'
    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)

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

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

    #FLOR_AM p
    ds = get_cycle_p_floram(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    label = 'p($\Lambda$)'
    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)

    #FLOR_AM nseed
    ds = get_cycle_nseed_floram(basin=basin, years=years)
    ds = ds.sel(life=seed_olife)
    ds = ds/ds.mclim.sum('month')
    label = 'N_SEED'
    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)

    #FLOR_AM nseed x p
    ds = get_cycle_nseedxp_floram(basin=basin, years=years)
    ds = ds.sel(life=seed_olife)
    ds = ds/ds.mclim.sum('month')
    label = 'N_SEED$\\times$p($\Lambda$)'
    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)

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

    #FLOR_AM+ p
    ds = get_cycle_p_floramplus(basin=basin, years=years)
    ds = ds/ds.mclim.sum('month')
    label = 'p($\Lambda$)'
    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)

    #FLOR_AM+ nseed
    ds = get_cycle_nseed_floramplus(basin=basin, years=years)
    ds = ds.sel(life=seed_olife)
    ds = ds/ds.mclim.sum('month')
    label = 'N_SEED'
    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)

    #FLOR_AM+ nseed x p
    ds = get_cycle_nseedxp_floramplus(basin=basin, years=years)
    ds = ds.sel(life=seed_olife)
    ds = ds/ds.mclim.sum('month')
    label = 'N_SEED$\\times$p($\Lambda$)'
    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)

    ax = axes[0, 0]
    ylabel = f'{basin} TC relative frequncy'
    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[0, 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('')
    ax.set_xlabel('')
    ax.set_xticks(range(1,13))
    ax.set_title('')
    ax.set_title('(b) HiRAM', loc=title_loc)

    ax = axes[1, 0]
    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_xlabel('month')
    ax.set_xticks(range(1,13))
    ax.set_title('')
    ax.set_title('(c) AM2.5', loc=title_loc)
    
    ax = axes[1, 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('')
    ax.set_xlabel('month')
    ax.set_xticks(range(1,13))
    ax.set_title('')
    ax.set_title('(d) AM2.5C360', 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()