#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Sep 29 14:12:57 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
maindir = '/tigress/wenchang/analysis/seedTC'
if maindir not in sys.path:
    sys.path.append(maindir)
from HiRAM_CTL1990.data_pTang_partial_cycle import get_cycle as get_cycle_ctl
from HiRAM_plus2K.data_pTang_oneCTL_cycle import get_cycle as get_cycle_plus2K
from HiRAM_2xCO2.data_pTang_oneCTL_cycle import get_cycle as get_cycle_2xCO2
from HiRAM_plus2K_2xCO2.data_pTang_oneCTL_cycle import get_cycle as get_cycle_plus2K_2xCO2


#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
def plot_cycle(ds, label=None, **kws):
    x = ds.month
    y1 = ds.mclim + ds.sem
    y2 = ds.mclim - ds.sem
    y = ds.mclim
    plt.fill_between(x, y1, y2, alpha=0.3, **kws)
    y.plot(marker='o', label=label, **kws)

def plot_cycle_change(ds, ds_ctl, label=None, **kws):
    da = (ds.mclim - ds_ctl.mclim)/ds_ctl.mclim * 100
    da = da.assign_attrs(long_name='probability change', units='%')
    da.plot(marker='o', label=label, **kws)

basin = 'NA'
ds_ctl = get_cycle_ctl(basin=basin)

ds_plus2K = get_cycle_plus2K(basin=basin)
ds_plus2K_ctlvshear = get_cycle_plus2K(basin=basin, fixterm='Vshear')
ds_plus2K_ctlchi = get_cycle_plus2K(basin=basin, fixterm='chi')
ds_plus2K_ctlPI = get_cycle_plus2K(basin=basin, fixterm='PI')

ds_2xCO2 = get_cycle_2xCO2(basin=basin)
ds_2xCO2_ctlvshear = get_cycle_2xCO2(basin=basin, fixterm='Vshear')
ds_2xCO2_ctlchi = get_cycle_2xCO2(basin=basin, fixterm='chi')
ds_2xCO2_ctlPI = get_cycle_2xCO2(basin=basin, fixterm='PI')

ds_plus2K_2xCO2 = get_cycle_plus2K_2xCO2(basin=basin)
ds_plus2K_2xCO2_ctlvshear = get_cycle_plus2K_2xCO2(basin=basin, fixterm='Vshear')
ds_plus2K_2xCO2_ctlchi = get_cycle_plus2K_2xCO2(basin=basin, fixterm='chi')
ds_plus2K_2xCO2_ctlPI = get_cycle_plus2K_2xCO2(basin=basin, fixterm='PI')
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    plt.close('all')
    #figname = __file__.replace('.py', f'_{basin}_{tt.today()}.png')
    fig, axes = plt.subplots(3, 1, figsize=(6,7), sharey=True)

    ax = axes[0]; plt.sca(ax)
    plot_cycle_change(ds_plus2K_2xCO2, ds_ctl, label='plus2K_2xCO2', color='k')
    plot_cycle_change(ds_plus2K_2xCO2_ctlvshear, ds_ctl, label='CTLclim Vshear')
    plot_cycle_change(ds_plus2K_2xCO2_ctlchi, ds_ctl, label='CTLclim $\chi$')
    plot_cycle_change(ds_plus2K_2xCO2_ctlPI, ds_ctl, label='CTLclim PI')

    plt.legend(loc='lower left')
    ax.grid('on')
    ax.set_xlabel('')

    ax = axes[1]; plt.sca(ax)
    plot_cycle_change(ds_plus2K, ds_ctl, label='plus2K', color='k')
    plot_cycle_change(ds_plus2K_ctlvshear, ds_ctl, label='CTLclim Vshear')
    plot_cycle_change(ds_plus2K_ctlchi, ds_ctl, label='CTLclim $\chi$')
    plot_cycle_change(ds_plus2K_ctlPI, ds_ctl, label='CTLclim PI')

    plt.legend(loc='lower left')
    ax.grid('on')
    ax.set_xlabel('')

    ax = axes[2]; plt.sca(ax)
    plot_cycle_change(ds_2xCO2, ds_ctl, label='2xCO2', color='k')
    plot_cycle_change(ds_2xCO2_ctlvshear, ds_ctl, label='CTLclim Vshear')
    plot_cycle_change(ds_2xCO2_ctlchi, ds_ctl, label='CTLclim $\chi$')
    plot_cycle_change(ds_2xCO2_ctlPI, ds_ctl, label='CTLclim PI')

    plt.legend(loc='lower left')
    ax.grid('on')

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