#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Feb  1 16:59:42 EST 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
import xfilter
from shared import spines_off
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#colorSpan = 'lightgray'
def plot_LME_TC(case, n_ens=None, year_start=850, year_end=2005, n_lowpass = 40, padtype='even', ens_color=None, ensm_color=None):
#     case = 'GHG'
#     n_ens = 3
    if n_ens > 5:
        add_legend = False
    else:
        add_legend = True
    if ensm_color is None: ensm_color = 'k'
    
    # ifile
    thisdir = os.path.dirname(__file__)
    if case in ['fullForcingOld',]:
        #ifile = f'/tigress/gvecchi/ANALYSIS/LMR_Paper1/LME_TC.nc'
        ifile = os.path.join(thisdir, 'data', f'LME_TC.nc') #wy 2021-02-05
    elif case in ['850forcing', '0850cntl',]:
        #ifile = f'/tigress/gvecchi/ANALYSIS/LMR_Paper1/LME_TC.{case}.nc'
        ifile = os.path.join(thisdir, 'data', f'LME_TC.{case}.fullForcingRef.nc') #wy 2021-02-05
    else:
        #ifile = f'/tigress/gvecchi/ANALYSIS/LMR_Paper1/LME_TC.{case}.{n_ens}ens.nc'
        ifile = os.path.join(thisdir, 'data', f'LME_TC.{case}.{n_ens}ens.fullForcingRef.nc') #wy 2021-02-05
    ds = xr.open_dataset(ifile)
    print('[loaded]:', ifile)

    # plot
    if case in ['850forcing', '0850cntl',]:
        ds.HU.filter.lowpass(1./n_lowpass,dim='year', padtype=padtype).plot()
    else:
        ds.HU.filter.lowpass(1./n_lowpass,dim='year', padtype=padtype).plot(hue='en', add_legend=add_legend, color=ens_color, lw=0.5, alpha=0.5)
        ds.HU.filter.lowpass(1./n_lowpass, dim='year', padtype=padtype).mean('en').plot(color=ensm_color, lw=2)

    #plt.axvspan(year_start, year_start+n_lowpass//2, color=colorSpan)
    #plt.axvspan(year_end-n_lowpass//2, year_end, color=colorSpan)
    plt.title(f'LME TC: {case}')
    plt.autoscale()
    spines_off( plt.gca() )
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    plot_LME_TC(case='fullForcing', n_ens=13, ensm_color='k', ens_color='gray')
    plt.xticks(np.arange(2000, 850, -100))
    plt.xlim(850, 2000)
    
    figname = __file__.replace('.py', f'.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        wysavefig(figname)
    else:
        print(f'figure not saved. to save figure: python {__file__} savefig')
    tt.check(f'**Done**')
    plt.show()