#!/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
from HiRAM_CTL1990.fig_scatter_ntc_vs_all import scatterplot
import xaddon
#
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)
"""
_ifile = 'fig_scatter_ntc_vs_all.nc'
case = 'HiRAM_CTL1990'
ifile = os.path.join(case, _ifile)
ds_ctl = xr.open_dataset(ifile)

case = 'HiRAM_plus2K'
ifile = os.path.join(case, _ifile)
ds_plus2K = xr.open_dataset(ifile)
    
case = 'HiRAM_2xCO2'
ifile = os.path.join(case, _ifile)
ds_2xCO2 = xr.open_dataset(ifile)

case = 'HiRAM_plus2K_2xCO2'
ifile = os.path.join(case, _ifile)
ds_plus2K_2xCO2 = xr.open_dataset(ifile)

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, axes = plt.subplots(4, 3, sharey=True, sharex='col', figsize=(6, 1.8*4))
    suptitle = None #'Obs. (IBTrACS TC)'
    #figname = __file__.replace('.py', f'_{tt.today()}.png')

    #obs
    case = 'CTL1990'
    ds = ds_ctl

    ax = axes[0, 0]
    scatterplot(ax=ax, xlabel=None, ylabel=f'{case} N_TC', x='p', y='ntc', data=ds, tag='(a)\n')

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

    ax = axes[0, 2]
    scatterplot(ax=ax, xlabel=None, x='nseedxp', y='ntc', data=ds, tag='(c)\n')

    case = 'plus2K'
    ds = ds_plus2K

    ax = axes[1, 0]
    scatterplot(ax=ax, ylabel=f'{case} N_TC', x='p', y='ntc', data=ds, tag='(d)\n')

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

    ax = axes[1, 2]
    scatterplot(ax=ax, x='nseedxp', y='ntc', data=ds, tag='(f)\n')

    case = '2xCO2'
    ds = ds_2xCO2

    ax = axes[2, 0]
    scatterplot(ax=ax, ylabel=f'{case} N_TC', x='p', y='ntc', data=ds, tag='(g)\n')

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

    ax = axes[2, 2]
    scatterplot(ax=ax, x='nseedxp', y='ntc', data=ds, tag='(i)\n')

    case = 'plus2K_2xCO2'
    ds = ds_plus2K_2xCO2

    ax = axes[3, 0]
    scatterplot(ax=ax, xlabel='p($\Lambda$)', ylabel=f'{case} N_TC', x='p', y='ntc', data=ds, tag='(j)\n')

    ax = axes[3, 1]
    scatterplot(ax=ax, xlabel='N_SEED', x='nseed', y='ntc', data=ds, tag='(k)\n')

    ax = axes[3, 2]
    scatterplot(ax=ax, xlabel='N_SEED$\\times$p($\Lambda$)', x='nseedxp', y='ntc', data=ds, tag='(l)\n')

    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()