#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed May 27 15:19:27 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 amipHadISST.data_ntc_cycle import get_cycle as get_cycle_amipHadISST
from ERA5.data_pTang_cycle import get_cycle as get_cycle_pTang_era5
from amipHadISST.data_pTang_cycle import get_cycle as get_cycle_pTang_amipHadISST
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
basin = 'NA'
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #figname = __file__.replace('.py', f'_{tt.today()}.png')
    fig, ax = plt.subplots()
    label = 'IBTrACS TC'
    ds = get_cycle_ibtracs(basin=basin)
    b1 = ax.bar(x='month', height='mclim', align='edge', width=-0.4, color='0.3', data=ds, label=label)
    ax.errorbar(data=ds.assign_coords(month=ds.month-0.2), x='month', y='mclim', yerr='cim', capsize=2, ls='none', color='k', label=None, alpha=1/2)
 
    label = 'HiRAM TC'
    ds = get_cycle_amipHadISST(basin=basin)
    b2 = ax.bar(x='month', height='mclim', align='edge', width=0.4, color='C0', data=ds, label=label)
    ax.errorbar(data=ds.assign_coords(month=ds.month+0.2), x='month', y='mclim', yerr='cim', capsize=2, ls='none', color='k', label=None, alpha=1/2)

    ax2 = ax.twinx()
    label = 'ERA5 p (right)'
    ds = get_cycle_pTang_era5(basin=basin)#.pipe(lambda x: x.assign_coords(month=x.month-0.2))
    ds['lower'] = ds['mclim'] - ds['cim']
    ds['upper'] = ds['mclim'] + ds['cim']
    ax2.fill_between(data=ds, x='month', y1='lower', y2='upper', alpha=1/4, color='k')
    l1 = ax2.plot('month', 'mclim', data=ds, color='k', label=label, marker='.')

    label = 'HiRAM p (right)'
    ds = get_cycle_pTang_amipHadISST(basin=basin)#.pipe(lambda x: x.assign_coords(month=x.month+0.2))
    ds['lower'] = ds['mclim'] - ds['cim']
    ds['upper'] = ds['mclim'] + ds['cim']
    ax2.fill_between(data=ds, x='month', y1='lower', y2='upper', alpha=1/4, color='C1')
    l2 = ax2.plot('month', 'mclim', data=ds, color='C1', label=label, marker='.', ls='--')

    ax.set_ylim(0, None)
    #ax.legend(loc='upper left', frameon=False)
    ax.set_xlabel('month')
    ax.set_ylabel('TC # monthly climatology')

    ax2.set_ylim(0, None)
    ax2.set_xticks(range(1, 13))
    lbs = [b1, b2] + l1 + l2
    labels = [lb.get_label() for lb in lbs]
    ax2.legend(lbs, labels, loc='upper left')
    ax2.set_xlabel('month')
    ax2.set_ylabel('p($\Lambda$)')

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