#!/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='TC genesis 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_2xCO2 = get_cycle_2xCO2(basin=basin)
ds_plus2K_2xCO2 = get_cycle_plus2K_2xCO2(basin=basin)
 
 
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(2, 1, figsize=(6,6))

    ax = axes[0]; plt.sca(ax)
    plot_cycle(ds_ctl, label='CTL', color='k')
    plot_cycle(ds_plus2K_2xCO2, label='plus2K_2xCO2')
    plot_cycle(ds_plus2K, label='plus2K')
    plot_cycle(ds_2xCO2, label='2xCO2')

    ax.set_ylabel('TC genesis probability')
    plt.legend()
    ax.grid('on')

    ax = axes[1]; plt.sca(ax)
    plot_cycle_change(ds_plus2K_2xCO2, ds_ctl, label='plus2K_2xCO2')
    plot_cycle_change(ds_plus2K, ds_ctl, label='plus2K')
    plot_cycle_change(ds_2xCO2, ds_ctl, label='2xCO2')

    plt.legend()
    ax.grid('on')
    ax.set_ylim(-50,10)

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