#!/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_rainy import scatterplot
from fig_scatter_ntc_vs_all_mbasins_v4 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_rainy.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)

#normalize by annual total value
func_norm = lambda x: x/x.sum('month')
ds_ctl = ds_ctl.pipe(func_norm)
ds_plus2K = ds_plus2K.pipe(func_norm)
ds_2xCO2 = ds_2xCO2.pipe(func_norm)
ds_plus2K_2xCO2 = ds_plus2K_2xCO2.pipe(func_norm)

# concatenate data
ds = xr.concat([ds_ctl, ds_plus2K, ds_2xCO2, ds_plus2K_2xCO2], dim=pd.Index(['CTL1990', 'plus2K', '2xCO2', 'plus2K_2xCO2'], name='exp'))
exps = ds.exp.values
ds = ds.stack(s=['exp', 'month'])

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

    #obs
    ax = axes[0]
    #scatterplot(ax=ax, xlabel=None, ylabel=f'{case} N_TC', x='p', y='ntc', data=ds, tag='(a)\n')
    scatterplot(ax=ax, xlabel='p($\Lambda$)', ylabel='N_TC', x='p', y='ntc', data=ds, title=None, tag='(a)', label=None, color='.99')
    for ee in exps:
        scatterplot(ax=ax, x='p', y='ntc', data=ds.sel(exp=ee), rg_on=False, label=ee)

    ax = axes[1]
    #scatterplot(ax=ax, xlabel=None, x='nseed', y='ntc', data=ds, tag='(b)\n')
    scatterplot(ax=ax, xlabel='N_SEED', x='nseed', y='ntc', data=ds, tag='(b)', color='.99', label=None)
    for ee in exps:
        scatterplot(ax=ax, x='nseed', y='ntc', data=ds.sel(exp=ee), rg_on=False, label=ee)

    ax = axes[2]
    #scatterplot(ax=ax, xlabel=None, x='nseedxp', y='ntc', data=ds, tag='(c)\n')
    scatterplot(ax=ax, xlabel='N_SEED$\\times$p($\Lambda$)', x='nseedxp', y='ntc', data=ds, tag='(c)', color='.99', label=None)
    for ee in exps:
        scatterplot(ax=ax, x='nseedxp', y='ntc', data=ds.sel(exp=ee), rg_on=False, label=ee)
    key = 'ntc'
    ymax = ds[key].max().item()
    f = 1.05 #scale factor
    ax.set_ylim(0, ymax*f)
    ax.set_xlim(0, ymax*f)
    
    plt.tight_layout(rect=[0,0,1,0.9])
    axes[2].legend(bbox_to_anchor=(1, 1), loc='lower right', frameon=False, ncol=4, borderpad=0)

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

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

    tt.check(f'**Done**')
    plt.show()
    plt.rcParams['figure.constrained_layout.use'] = True