#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Mar 11 14:56:49 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, matplotlib.pyplot as plt
#more imports
from misc.cim import cim
from shared import spines_off
#from fig_lines_tc_all import ds_lmr as ds_lmr2019, ds_lmr2018, lowpass
from fig_lines_tc_all_poisson_HUcount_tfrMDRTROP18702000 import ds_lmr2019tf as ds_lmr2019, ds_lmr2018tf as ds_lmr2018, lowpass
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
year0 = 850
lmr_window = 40
alpha_fill = 0.1
lw = 2

#if True: 
    #from wyconfig import *
def wyplot(ax=None):
    if ax is None:
        fig, ax = plt.subplots(figsize=(8,3))
    qq = [0.25, 0.75]#lower and upper quartiles
    #lmr2018 MDR-TROP
    name = 'MDR$-$TROP'
    label = 'LMR2.0 ' + name
    ds_lmr = ds_lmr2018
    da = (ds_lmr['MDR'] - ds_lmr['TROP']).mean('mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    """
    err = (ds_lmr['MDR'] - ds_lmr['TROP']).pipe(cim, dim='mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    ax.fill_between(da.year, da-err, da+err, alpha=alpha_fill, color='C0')
    ci = (ds_lmr['MDR'] - ds_lmr['TROP']).pipe(lowpass) \
        .quantile(qq, dim='mcrun') \
        .sel(year=slice(year0,None))
    ax.fill_between(ci.year, ci[0], ci[1], alpha=alpha_fill, color='C0')
    """
    stde = (ds_lmr['MDR'] - ds_lmr['TROP']).pipe(lowpass) \
        .std('mcrun') \
        .sel(year=slice(year0,None))
    ci = da - stde, da + stde
    ax.fill_between(da.year, ci[0], ci[1], alpha=alpha_fill, color='C0')
    da.plot(ax=ax, label=label, lw=lw, color='C0')
    #lmr2019 MDR-TROP
    name = 'MDR$-$TROP'
    label = 'LMR2.1 ' + name
    ds_lmr = ds_lmr2019
    da = (ds_lmr['MDR'] - ds_lmr['TROP']).mean('mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    """
    err = (ds_lmr['MDR'] - ds_lmr['TROP']).pipe(cim, dim='mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    ax.fill_between(da.year, da-err, da+err, alpha=alpha_fill, color='C1')
    ci = (ds_lmr['MDR'] - ds_lmr['TROP']).pipe(lowpass) \
        .quantile(qq, dim='mcrun') \
        .sel(year=slice(year0,None))
    ax.fill_between(ci.year, ci[0], ci[1], alpha=alpha_fill, color='C1')
    """
    stde = (ds_lmr['MDR'] - ds_lmr['TROP']).pipe(lowpass) \
        .std('mcrun') \
        .sel(year=slice(year0,None))
    ci = da - stde, da + stde
    ax.fill_between(da.year, ci[0], ci[1], alpha=alpha_fill, color='C1')
    da.plot(ax=ax, label=label, lw=lw, color='C1')
    #lmr2018 MDR
    name = 'MDR'
    label = 'LMR2.0 ' + name
    ds_lmr = ds_lmr2018
    da = ds_lmr[name].mean('mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    """
    err = ds_lmr[name].pipe(cim, dim='mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    ax.fill_between(da.year, da-err, da+err, alpha=alpha_fill, color='C0')
    ci = ds_lmr[name].pipe(lowpass) \
        .quantile(qq, dim='mcrun') \
        .sel(year=slice(year0,None))
    ax.fill_between(ci.year, ci[0], ci[1], alpha=alpha_fill, color='C0')
    """
    stde = ds_lmr[name].pipe(lowpass) \
        .std('mcrun') \
        .sel(year=slice(year0,None))
    ci = da - stde, da + stde
    ax.fill_between(da.year, ci[0], ci[1], alpha=alpha_fill, color='C0')
    da.plot(ax=ax, label=label, lw=lw, color='C0', ls='--')
    #lmr2019 MDR
    name = 'MDR'
    label = 'LMR2.1 ' + name
    ds_lmr = ds_lmr2019
    da = ds_lmr[name].mean('mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    """
    err = ds_lmr[name].pipe(cim, dim='mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    ax.fill_between(da.year, da-err, da+err, alpha=alpha_fill, color='C1')
    ci = ds_lmr[name].pipe(lowpass) \
        .quantile(qq, dim='mcrun') \
        .sel(year=slice(year0,None))
    ax.fill_between(ci.year, ci[0], ci[1], alpha=alpha_fill, color='C1')
    """
    stde = ds_lmr[name].pipe(lowpass) \
        .std('mcrun') \
        .sel(year=slice(year0,None))
    ci = da - stde, da + stde
    ax.fill_between(da.year, ci[0], ci[1], alpha=alpha_fill, color='C1')
    da.plot(ax=ax, label=label, lw=lw, color='C1', ls='--')
    #lmr2018 TROP
    name = 'TROP'
    label = 'LMR2.0 ' + name
    ds_lmr = ds_lmr2018
    da = ds_lmr[name].mean('mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    """
    err = ds_lmr[name].pipe(cim, dim='mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    ax.fill_between(da.year, da-err, da+err, alpha=alpha_fill, color='C0')
    ci = ds_lmr[name].pipe(lowpass) \
        .quantile(qq, dim='mcrun') \
        .sel(year=slice(year0,None))
    ax.fill_between(ci.year, ci[0], ci[1], alpha=alpha_fill, color='C0')
    """
    stde = ds_lmr[name].pipe(lowpass) \
        .std('mcrun') \
        .sel(year=slice(year0,None))
    ci = da - stde, da + stde
    ax.fill_between(da.year, ci[0], ci[1], alpha=alpha_fill, color='C0')
    da.plot(ax=ax, label=label, lw=lw, color='C0', ls=':')
    #lmr2019 TROP
    name = 'TROP'
    label = 'LMR2.1 ' + name
    ds_lmr = ds_lmr2019
    da = ds_lmr[name].mean('mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    """
    err = ds_lmr[name].pipe(cim, dim='mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    ax.fill_between(da.year, da-err, da+err, alpha=alpha_fill, color='C1')
    ci = ds_lmr[name].pipe(lowpass) \
        .quantile(qq, dim='mcrun') \
        .sel(year=slice(year0,None))
    ax.fill_between(ci.year, ci[0], ci[1], alpha=alpha_fill, color='C1')
    """
    stde = ds_lmr[name].pipe(lowpass) \
        .std('mcrun') \
        .sel(year=slice(year0,None))
    ci = da - stde, da + stde
    ax.fill_between(da.year, ci[0], ci[1], alpha=alpha_fill, color='C1')
    da.plot(ax=ax, label=label, lw=lw, color='C1', ls=':')
    
    #
    ax.axhline(0, color='gray', ls='--')
    spines_off(ax)
    ax.set_title(f'MDR and tropical mean SST anomomalies (lowpass_cutoff={lmr_window})', loc='left')
    ax.set_ylabel('K')
    ax.set_ylim(-0.6, 0.6)
    ax.set_xticks(range(900, 2000+1, 100))
    ax.set_xlim(year0, 2000)
    ax.legend(ncol=3, frameon=False, loc='upper right')

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)
    tt.check(f'**Done**')
    plt.show()