#!/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
from misc.cim import cim
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
def wyplot(ax=None, legend_on=True, legend_ncol=3):
    if ax is None:
        fig, ax = plt.subplots(figsize=(8, 3))
    n_window = 40
    padtype = 'even'
    cases = ['fullForcing', 'GHG', 'VOLC_GRA', 
             'LULC_HurttPongratz', 'SSI_VSK_L', 'ORBITAL', 'OZONE_AER']#, '850forcing']#, '0850cntl']
    #case info: https://www.cesm.ucar.edu/experiments/cesm1.1/LM/
    for case in cases:
        color = 'k' if case == 'fullForcing' else None
        lw = 2 if case == 'fullForcing' else 1.5
        alpha = 1 if case == 'fullForcing' else 0.5
        ifile = os.path.join(os.path.dirname(__file__), 'data', f'LME_TC.{case}.*fullForcingRef.nc')
        ds = xr.open_mfdataset(ifile, combine='by_coords')
        print('[loaded]:', ifile)
        if 'en' in ds.dims:
            da = ds.TROP.mean('en').filter.lowpass(1./n_window, dim='year', padtype=padtype)
            n_ens = ds.en.size
        else:
            da = ds.TROP.filter.lowpass(1./n_window, dim='year', padtype=padtype)
            n_ens = None
        if n_ens:
            label = f'{case}({n_ens})'
        else:
            label = case
        da.plot(ax=ax, color=color, label=label, alpha=alpha, lw=lw)
        if False and case == 'fullForcing':
            da_cim = ds.TROP.filter.lowpass(1./n_window, dim='year', padtype=padtype).pipe(cim, confidence=0.95)
            ax.fill_between(da.year, da-da_cim, da+da_cim, color='gray', alpha=0.3)

    ifile1 = os.path.join(os.path.dirname(__file__), 'data', f'LME_TC.0850cntl.fullForcingRef.nc')
    ifile2 = os.path.join(os.path.dirname(__file__), 'data', f'LME_TC.850forcing.fullForcingRef.nc')
    #da_ctl = ( xr.open_dataset(ifile1).TROP.mean('year') + xr.open_dataset(ifile2).TROP.mean('year') )/2 # average between two ctl runs
    #da_ctl_std = ( xr.open_dataset(ifile1).TROP.filter.lowpass(1./n_window, dim='year', padtype=padtype).std('year') + xr.open_dataset(ifile2).TROP.filter.lowpass(1./n_window, dim='year', padtype=padtype).std('year') )/2 # average std between two ctl runs
    da_ctl = xr.open_dataset(ifile1).TROP.mean('year') 
    #da_ctl_std = xr.open_dataset(ifile1).TROP.filter.lowpass(1./n_window, dim='year', padtype=padtype).std('year') 
    #lower = da_ctl - da_ctl_std
    #upper = da_ctl + da_ctl_std
    da_ = xr.open_dataset(ifile1).TROP.filter.lowpass(1./n_window, dim='year', padtype=padtype)
    lower = da_.quantile(0.025, dim='year')
    upper = da_.quantile(0.975, dim='year')
    #ax.axhline(da_ctl, color='gray', ls='--', label='0850cntl mean and mean$\pm$STD')
    ax.axhline(da_ctl, color='gray', ls='--', label='0850cntl mean and 95% CI')
    ax.axhline(lower, color='gray', ls='--')
    ax.axhline(upper, color='gray', ls='--')
    #plt.axhspan(da_ctl-da_ctl_std, da_ctl+da_ctl_std, color='gray', alpha=0.3)

    if legend_on:
        ax.legend(ncol=legend_ncol, frameon=False)
    ax.set_xticks(np.arange(2000, 850, -100))
    ax.set_xlim(850, 2000)
    ax.set_ylim(24.8, 26.4)
    spines_off(ax)
    ax.set_ylabel('Tropical mean SST [$^\circ$C]')
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    wyplot()

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