#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Jun  4 11:50:59 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
import xaddon
from ERA5.fig_scatter_ntc_vs_all_rainy import get_scatter_data as get_scatter_data_era5
from IBTrACS.data_ntc_cycle import get_cycle as cycle_ntc
from amipHadISST.fig_scatter_ntc_vs_all_rainy import get_scatter_data as get_scatter_data_hiram, scatterplot
from AM2p5.fig_scatter_ntc_vs_all_rainy import get_scatter_data as get_scatter_data_floram
from AM2p5C360.fig_scatter_ntc_vs_all_rainy import get_scatter_data as get_scatter_data_floramplus
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
basin = 'NA'
years = slice('1980', '2018')
ds_obs = get_scatter_data_era5(basin=basin, years=years)
ds_obs['ntc_era5'] = ds_obs['ntc']
ds_obs['ntc'] = cycle_ntc(basin=basin, years=years)['mclim']
ds_hiram = get_scatter_data_hiram(basin=basin, years=years)
ds_floram = get_scatter_data_floram(basin=basin, years=years)
ds_floramplus = get_scatter_data_floramplus(basin=basin, years=years)
    
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, axes = plt.subplots(4, 3, sharey=True, sharex='col', figsize=(7.5,6.5))
    suptitle = None #'Obs. (IBTrACS TC)'
    #figname = __file__.replace('.py', f'_{tt.today()}.png')

    #obs
    ds = ds_obs

    ax = axes[0, 0]
    #scatterplot(ax=ax, xlabel='ERA5 p($\Lambda$)', ylabel='IBTrACS N_TC', x='p', y='ntc', data=ds)
    #scatterplot(ax=ax, xlabel=None, ylabel='N_TC', x='p', y='ntc', data=ds, title='Obs.', tag='(a)')
    scatterplot(ax=ax, xlabel=None, ylabel='N_TC', x='p', y='ntc', data=ds, title=None, tag='(a)')

    ax = axes[0, 1]
    #scatterplot(ax=ax, xlabel='N_SEED', x='nseed', y='ntc', data=ds)
    scatterplot(ax=ax, xlabel=None, x='nseed', y='ntc', data=ds, tag='(b)')

    ax = axes[0, 2]
    #scatterplot(ax=ax, xlabel='N_SEED$\\times$p($\Lambda$)', x='nseedxp', y='ntc', data=ds)
    scatterplot(ax=ax, xlabel=None, x='nseedxp', y='ntc', data=ds, tag='(c)')
    ax.text(0.99, 0.01, 'Obs.', transform=ax.transAxes, ha='right', va='bottom')

    ds = ds_hiram
    ax = axes[1, 0]
    #scatterplot(ax=ax, xlabel='p($\Lambda$)', ylabel='HiRAM N_TC', x='p', y='ntc', data=ds)
    #scatterplot(ax=ax, xlabel=None, ylabel='N_TC', x='p', y='ntc', data=ds, title='HiRAM', tag='(d)')
    scatterplot(ax=ax, xlabel=None, ylabel='N_TC', x='p', y='ntc', data=ds, title=None, tag='(d)')

    ax = axes[1, 1]
    scatterplot(ax=ax, xlabel=None, x='nseed', y='ntc', data=ds, tag='(e)')

    ax = axes[1, 2]
    scatterplot(ax=ax, xlabel=None, x='nseedxp', y='ntc', data=ds, tag='(f)')
    ax.text(0.99, 0.01, 'HiRAM', transform=ax.transAxes, ha='right', va='bottom')

    ds = ds_floram
    ax = axes[2, 0]
    #scatterplot(ax=ax, xlabel=None, ylabel='N_TC', x='p', y='ntc', data=ds, title='AM2.5', tag='(g)')
    scatterplot(ax=ax, xlabel=None, ylabel='N_TC', x='p', y='ntc', data=ds, title=None, tag='(g)')

    ax = axes[2, 1]
    scatterplot(ax=ax, xlabel=None, x='nseed', y='ntc', data=ds, tag='(h)')

    ax = axes[2, 2]
    scatterplot(ax=ax, xlabel=None, x='nseedxp', y='ntc', data=ds, tag='(i)')
    ax.text(0.99, 0.01, 'AM2.5', transform=ax.transAxes, ha='right', va='bottom')

    ds = ds_floramplus
    ax = axes[3, 0]
    #scatterplot(ax=ax, xlabel='p($\Lambda$)', ylabel='N_TC', x='p', y='ntc', data=ds, title='AM2.5C360', tag='(j)')
    scatterplot(ax=ax, xlabel='p($\Lambda$)', ylabel='N_TC', x='p', y='ntc', data=ds, title=None, tag='(j)')
    ax.set_xlim(0, 0.18) # key = 'p'; max( ds_obs[key].max(), ds_hiram[key].max(), ds_floram[key].max(), ds_floramplus[key].max())

    ax = axes[3, 1]
    scatterplot(ax=ax, xlabel='N_SEED', x='nseed', y='ntc', data=ds, tag='(k)')
    ax.set_xlim(0, 31) # key = 'nseed'; max( ds_obs[key].max(), ds_hiram[key].max(), ds_floram[key].max(), ds_floramplus[key].max())

    ax = axes[3, 2]
    scatterplot(ax=ax, xlabel='N_SEED$\\times$p($\Lambda$)', x='nseedxp', y='ntc', data=ds, tag='(l)')
    ax.set_xlim(0, 3.6) # key = 'nseedxp'; max( ds_obs[key].max(), ds_hiram[key].max(), ds_floram[key].max(), ds_floramplus[key].max())
    ax.text(0.99, 0.01, 'AM2.5C360', transform=ax.transAxes, ha='right', va='bottom')



    if suptitle is not None:
        fig.suptitle(suptitle)

    if len(sys.argv)>1 and sys.argv[1]=='savefig':
        figname = __file__.replace('.py', '.png')
        wysavefig(figname)

    tt.check(f'**Done**')
    plt.show()