#!/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
from AM2p5.data_pTang_cycle import get_cycle as get_cycle_pTang_am2p5
from AM2p5C360.data_pTang_cycle import get_cycle as get_cycle_pTang_am2p5c360
from ERA5.data_pTang_partial_cycle import get_cycle as get_cycle_pTangs_era5
from amipHadISST.data_pTang_partial_cycle import get_cycle as get_cycle_pTangs_amipHadISST
from AM2p5.data_pTang_partial_cycle import get_cycle as get_cycle_pTangs_am2p5
from AM2p5C360.data_pTang_partial_cycle import get_cycle as get_cycle_pTangs_am2p5c360
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
basin = 'NA'
years = slice('1980', '2018')

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    plt.close()
    #figname = __file__.replace('.py', f'_{tt.today()}.png')
    figsize = (8,5)
    fig, axes = plt.subplots(2, 2, figsize=figsize, sharex=True, sharey=True)

    ax = axes[0, 0]
    dc = {v:get_cycle_pTangs_era5(basin=basin, years=years, fixterm=v) for v in [None, 'Vshear', 'chi', 'PI']}
    dc[None].mclim.plot(label='ERA5 full', color='k', marker='o', ax=ax)
    dc['PI'].mclim.plot(label='const $u_{PI}$', color='C1', marker='o', ax=ax)
    dc['chi'].mclim.plot(label='const $\chi$', color='C2', marker='o', ax=ax)
    dc['Vshear'].mclim.plot(label='const $u_{shear}$', color='C3', marker='o', ax=ax)

    ax.legend()
    #ax.yaxis.tick_right()
    #ax.yaxis.set_label_position("right")
    ax.set_xlabel('')
    ax.set_ylabel('p')
    ax.set_ylim(ax_right.get_ylim())
    ax.text(1-0.01, 1-0.02, '(a)', ha='right', va='top', transform=ax.transAxes, fontsize='large')

    ax = axes[0, 1]
    dc = {v:get_cycle_pTangs_amipHadISST(basin=basin, years=years, fixterm=v) for v in [None, 'Vshear', 'chi', 'PI']}
    dc[None].mclim.plot(label='HiRAM full', color='C0', marker='o', ax=ax)
    dc['PI'].mclim.plot(label='const $u_{PI}$', color='C1', marker='o', ax=ax)
    dc['chi'].mclim.plot(label='const $\chi$', color='C2', marker='o', ax=ax)
    dc['Vshear'].mclim.plot(label='const $u_{shear}$', color='C3', marker='o', ax=ax)

    ax.legend()
    #ax.yaxis.tick_right()
    #ax.yaxis.set_label_position("right")
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_ylim(ax_right.get_ylim())
    ax.text(1-0.01, 1-0.02, '(b)', ha='right', va='top', transform=ax.transAxes, fontsize='large')

    ax = axes[1, 0]
    dc = {v:get_cycle_pTangs_am2p5(basin=basin, years=years, fixterm=v) for v in [None, 'Vshear', 'chi', 'PI']}
    dc[None].mclim.plot(label='AM2.5 full', color='C0', marker='o', ax=ax)
    dc['PI'].mclim.plot(label='const $u_{PI}$', color='C1', marker='o', ax=ax)
    dc['chi'].mclim.plot(label='const $\chi$', color='C2', marker='o', ax=ax)
    dc['Vshear'].mclim.plot(label='const $u_{shear}$', color='C3', marker='o', ax=ax)

    ax.legend()
    #ax.yaxis.tick_right()
    #ax.yaxis.set_label_position("right")
    ax.set_xlabel('month')
    ax.set_ylabel('p')
    ax.set_ylim(ax_right.get_ylim())
    ax.text(1-0.01, 1-0.02, '(c)', ha='right', va='top', transform=ax.transAxes, fontsize='large')

    ax = axes[1, 1]
    dc = {v:get_cycle_pTangs_am2p5c360(basin=basin, years=years, fixterm=v) for v in [None, 'Vshear', 'chi', 'PI']}
    dc[None].mclim.plot(label='AM2.5C360 full', color='C0', marker='o', ax=ax)
    dc['PI'].mclim.plot(label='const $u_{PI}$', color='C1', marker='o', ax=ax)
    dc['chi'].mclim.plot(label='const $\chi$', color='C2', marker='o', ax=ax)
    dc['Vshear'].mclim.plot(label='const $u_{shear}$', color='C3', marker='o', ax=ax)

    ax.legend()
    #ax.yaxis.tick_right()
    #ax.yaxis.set_label_position("right")
    #ax.set_xlabel('month')
    ax.set_ylabel('')
    ax.set_ylim(ax_right.get_ylim())
    ax.text(1-0.01, 1-0.02, '(d)', ha='right', va='top', transform=ax.transAxes, fontsize='large')
    ax.set_xticks(range(1,13))

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