#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Jan 29 15:18:06 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
#import matplotlib.pyplot as plt
#more imports
import xfilter
from misc.cim import cim
from data.data_adjusted_hurricanes import ds as ds_ah
from shared import emulate_hu
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
# params
thisdir = os.path.dirname(__file__)
# data
# tc liz
#ifile = '/tigress/gvecchi/ANALYSIS/LMR_Paper1/normTCcounts_v0719.txt'
ifile = '/tigress/gvecchi/ANALYSIS/LMR_Paper1/normTCcounts_v1219.txt'
tc_liz = pd.read_csv(ifile, sep='\s+', index_col=0)
# tc jk
#ifile = '/tigress/gvecchi/ANALYSIS/LMR_Paper1/jknife_v0719.txt'
ifile = '/tigress/gvecchi/ANALYSIS/LMR_Paper1/jknife_v1219.txt'
tc_jk = pd.read_csv(ifile, sep='\s+', index_col=0)

# tc lmr 2018
#ifile = '/tigress/gvecchi/DATA/LMR_2018/tc_rec.nc'
#lmr = ifile.split('/')[-2]
#ds = xr.open_dataset(ifile, decode_times=False)
#ds_lmr = ds.rename({'TIME': 'year', 'MCRUN': 'mcrun'}) \
#    .assign_coords(year=ds.TIME.values.astype('int').astype('float'))
ifile = os.path.join(os.path.dirname(__file__), 'data', 'LMR2018_TC.nc')
lmr = ifile.split('/')[-1].split('_')[0]
ds = xr.open_dataset(ifile)
ds_lmr = ds.rename(MCrun='mcrun')
#WY2021-02-17: use yearsRef of 1979-2000 (equivalent to 1982-2005 for HadISST)
hu = emulate_hu(ds_lmr, 'MDR', 'TROP', yearsRef=slice(1979,2000))
ds_lmr['HU'] = hu

 
year0 = 850
lmr_window = 40#61
alpha_fill = 0.2
# set top/righg spines off
def spines_off(ax):
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
def lowpass(da):
    '''lowpass'''
    #return da.filter.lowpass(1.0/lmr_window)
    return da.filter.lowpass(1.0/lmr_window, dim='year', padtype='even') #wy 2021-02-05
    
def plot_lines_liz(ax, projected=True):
    if projected:# convert from normalized TC to none-normalized TC
        normTCcountsMean = tc_liz['normTCcounts'][1901:2000].mean()
        normTCcountsSTD = tc_liz['normTCcounts'][1901:2000].std()
        ahMean = ds_ah['AHlp'].sel(year=slice(1901,2000)).mean().item()
        ahSTD = ds_ah['AHlp'].sel(year=slice(1901,2000)).std().item()
        tc_liz_ = (tc_liz - normTCcountsMean)/normTCcountsSTD * ahSTD + ahMean
        tc_jk_ = (tc_jk - normTCcountsMean)/normTCcountsSTD * ahSTD + ahMean
    else:
        normTCcountsMean = tc_liz['normTCcounts'][1901:2000].mean()
        normTCcountsSTD = tc_liz['normTCcounts'][1901:2000].std()
        #modern hurricane
        ahMean = ds_ah['AHlp'].sel(year=slice(1901,2000)).mean().item()
        ahSTD = ds_ah['AHlp'].sel(year=slice(1901,2000)).std().item()
        hu_ = (ds_ah['AHlp'] - ahMean)/ahSTD * normTCcountsSTD + normTCcountsMean
        #modern major hurricane
        amhMean = ds_ah['AMHlp'].sel(year=slice(1901,2000)).mean().item()
        amhSTD = ds_ah['AMHlp'].sel(year=slice(1901,2000)).std().item()
        mhu_ = (ds_ah['AMHlp'] - amhMean)/amhSTD * normTCcountsSTD + normTCcountsMean
        tc_liz_ = tc_liz
        tc_jk_ = tc_jk
    ll, uu = tc_liz_['smoothlower'][year0:None], tc_liz_['smoothupper'][year0:None]
    ax.fill_between(ll.index, ll, uu, alpha=alpha_fill, color='k')
    tc_liz_['normTCcounts'][year0:None].plot(ax=ax, color='k', lw=2, label='reconstructed HU')
    #tc_jk[year0:None].plot(ax=ax, lw=1, alpha=0.5)
    #sort the sites from north to south
    sites = ['noNewEng', 'noMidAtl', 'noSoutheast', 'noGulf_ChoBay', 'noGulf_AppBay',
        'noLIBHBahamas', 'noAbaco', 'noAndrosAM4', 'noVieqPR', 'noBelizeLRBH']
    tc_jk_[year0:None][sites].plot(ax=ax, lw=1, alpha=0.5)
    if projected:
        ds_ah['AHlp'].plot(ax=ax, color='k', ls='--', lw=2, label='Adjusted Hurricanes')
    else:
        hu_.plot(ax=ax, color='gray', ls='--', lw=2, label='modern recorded HU')
        mhu_.plot(ax=ax, color='k', ls='--', lw=2, label='modern recorded MH')
        
    ax.legend(ncol=4, frameon=False, loc='upper left')
    spines_off(ax)
    ax.set_title('TC proxy index from multi-site sediments', loc='left')
    if projected:
        ax.set_ylim(None, 10)
        ax.set_ylabel('#')
    else:
        ax.set_ylim(None, 0.07)
        ax.set_ylabel('normalized HU #')

def plot_hu_lmr(ax, liz_hu_on=False, modern_hu_on=False):
    name = 'HU'
    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=0.2, label='%95 CI')
    da.plot(ax=ax, label='SST-emulated HU', lw=2)
    for r in ds_lmr.mcrun:
        da = ds_lmr[name].sel(mcrun=r) \
            .pipe(lowpass) \
            .sel(year=slice(year0,None))
        da.plot(ax=ax, lw=1, color='C0', alpha=0.2)
    if liz_hu_on:
        lmrMean = ds_lmr[name].mean('mcrun').pipe(lowpass).sel(year=slice(1901,2000)).mean().item()
        lmrSTD = ds_lmr[name].mean('mcrun').pipe(lowpass).sel(year=slice(1901,2000)).std().item()
        func_norm = lambda x: (x - x['normTCcounts'][1901:2000].mean())/x['normTCcounts'][1901:2000].std() * lmrSTD + lmrMean
        tc_liz.pipe(func_norm)['normTCcounts'].plot(ax=ax, color='k', ls='-', lw=2, label='reconstructed HU')
    if modern_hu_on:
        lmrMean = ds_lmr[name].mean('mcrun').pipe(lowpass).sel(year=slice(1901,2000)).mean()
        lmrSTD = ds_lmr[name].mean('mcrun').pipe(lowpass).sel(year=slice(1901,2000)).std()
        func_norm = lambda x: (x - x.sel(year=slice(1901,2000)).mean())/x.sel(year=slice(1901,2000)).std() * lmrSTD + lmrMean
        ds_ah['AHlp'].pipe(func_norm).plot(ax=ax, color='gray', ls='--', lw=2, label='modern recorded HU')
        ds_ah['AMHlp'].pipe(func_norm).plot(ax=ax, color='k', ls='--', lw=2, label='modern recorded MH')
    if liz_hu_on or modern_hu_on:
        ax.legend(frameon=False, loc='upper left', ncol=2)
    #if projected:
    #    ds_ah['AHlp'].plot(ax=ax, color='k', ls='--', lw=2, label='Adjusted Hurricanes')
    #ax.legend(loc='upper left', ncol=2)
    spines_off(ax)
    ax.set_ylim(None, 7.5)
    ax.set_ylabel('HU #')
    ax.set_title(f'TC count from statistical model using {lmr} SST (lowpass_cutoff={lmr_window})', loc='left')

def plot_mdr_lmr(ax):
    #MDR - TROP
    name = 'MDR$-$TROP'
    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=0.4)
    da.plot(ax=ax, label=name)
    for r in ds_lmr.mcrun:
        da = (ds_lmr['MDR'] - ds_lmr['TROP']).sel(mcrun=r) \
            .pipe(lowpass) \
            .sel(year=slice(year0,None))
        da.plot(ax=ax, lw=1, color='C0', alpha=0.2)

    name = 'TROP'
    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=0.4)
    da.plot(ax=ax, label=name)
    for r in ds_lmr.mcrun:
        da = ds_lmr[name].sel(mcrun=r) \
            .pipe(lowpass) \
            .sel(year=slice(year0,None))
        da.plot(ax=ax, lw=1, color='C1', alpha=0.2)
    
    name = 'MDR'
    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=0.4)
    da.plot(ax=ax, label=name)
    for r in ds_lmr.mcrun:
        da = ds_lmr[name].sel(mcrun=r) \
            .pipe(lowpass) \
            .sel(year=slice(year0,None))
        da.plot(ax=ax, lw=1, color='C2', alpha=0.2)

    ax.legend(ncol=3, frameon=False)
    spines_off(ax)
    ax.set_title(f'MDR and tropical mean SSTs (lowpass_cutoff={lmr_window})', loc='left')
    ax.set_ylabel('K')
    ax.set_ylim(-0.4, 0.4)
    
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    projected = False # convert from normalized TC to none-normalized TC
    major_hurricane = True
    if projected:
        normalized = False
    else:
        normalized = True#modern hurricane normalized to proxy
    modern_hurricane_on = True
    liz_add_on_lmr = True
    
    # fig
    fig, axes = plt.subplots(3, 1, figsize=(8, 6+2), sharex=True, dpi=90)
    # 
    ax = axes[0]
    if projected:
        normTCcountsMean = tc_liz['normTCcounts'][1901:2000].mean()
        normTCcountsSTD = tc_liz['normTCcounts'][1901:2000].std()
        if major_hurricane:
            ahMean = ds_ah['AMHlp'].sel(year=slice(1901,2000)).mean().item()
            ahSTD = ds_ah['AMHlp'].sel(year=slice(1901,2000)).std().item()
        else:
            ahMean = ds_ah['AHlp'].sel(year=slice(1901,2000)).mean().item()
            ahSTD = ds_ah['AHlp'].sel(year=slice(1901,2000)).std().item()
        tc_liz = (tc_liz - normTCcountsMean)/normTCcountsSTD * ahSTD + ahMean
        tc_jk = (tc_jk - normTCcountsMean)/normTCcountsSTD * ahSTD + ahMean
    if normalized:
        normTCcountsMean = tc_liz['normTCcounts'][1901:2000].mean()
        normTCcountsSTD = tc_liz['normTCcounts'][1901:2000].std()
        #hurricane
        ahMean = ds_ah['AHlp'].sel(year=slice(1901,2000)).mean().item()
        ahSTD = ds_ah['AHlp'].sel(year=slice(1901,2000)).std().item()
        hu = (ds_ah['AHlp'] - ahMean)/ahSTD * normTCcountsSTD + normTCcountsMean
        #major hurricane
        amhMean = ds_ah['AMHlp'].sel(year=slice(1901,2000)).mean().item()
        amhSTD = ds_ah['AMHlp'].sel(year=slice(1901,2000)).std().item()
        mhu = (ds_ah['AMHlp'] - amhMean)/amhSTD * normTCcountsSTD + normTCcountsMean
    ll, uu = tc_liz['smoothlower'][year0:None], tc_liz['smoothupper'][year0:None]
    ax.fill_between(ll.index, ll, uu, alpha=alpha_fill, color='k')
    tc_liz['normTCcounts'][year0:None].plot(ax=ax, color='k', lw=2, label='reconstructed HU')
    #tc_jk[year0:None].plot(ax=ax, lw=1, alpha=0.5)
    #sort the sites from north to south
    sites = ['noNewEng', 'noMidAtl', 'noSoutheast', 'noGulf_ChoBay', 'noGulf_AppBay',
        'noLIBHBahamas', 'noAbaco', 'noAndrosAM4', 'noVieqPR', 'noBelizeLRBH']
    tc_jk[year0:None][sites].plot(ax=ax, lw=1, alpha=0.5)
    if projected:
        if major_hurricane:
            ds_ah['AMHlp'].plot(ax=ax, color='k', ls='--', lw=2, label='Adjusted Hurricanes')
        else:
            ds_ah['AHlp'].plot(ax=ax, color='k', ls='--', lw=2, label='Adjusted Hurricanes')
    if normalized:
        hu.plot(ax=ax, color='gray', ls='--', lw=2, label='recorded HU')
        mhu.plot(ax=ax, color='k', ls='--', lw=2, label='recorded MH')
    ax.legend(ncol=4, frameon=False, loc='upper left')
    # ax.set_ylim(0, 0.07)
    spines_off(ax)
    ax.set_title('TC proxy index from multi-site sediments', loc='left')
    if projected:
        if major_hurricane:
            pass
        else:
            ax.set_ylim(None, 10)
        ax.set_ylabel('#')
    else:
        ax.set_ylim(None, 0.07)
        ax.set_ylabel('normalized HU #')
    
    # ax.spines['bottom'].set_visible(False)
    
    #
    ax = axes[1]
    name = 'HU'
    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=0.2, label='%95 CI')
    da.plot(ax=ax, label='SST-emulated HU', lw=2)
    for r in ds_lmr.mcrun:
        da = ds_lmr[name].sel(mcrun=r) \
            .pipe(lowpass) \
            .sel(year=slice(year0,None))
        da.plot(ax=ax, lw=1, color='C0', alpha=0.2)
    #if projected:
    #    ds_ah['AHlp'].plot(ax=ax, color='k', ls='--', lw=2, label='Adjusted Hurricanes')
    if liz_add_on_lmr:
        lmrMean = ds_lmr[name].mean('mcrun').pipe(lowpass).sel(year=slice(1901,2000)).mean().item()
        lmrSTD = ds_lmr[name].mean('mcrun').pipe(lowpass).sel(year=slice(1901,2000)).std().item()
        func_norm = lambda x: (x - x['normTCcounts'][1901:2000].mean())/x['normTCcounts'][1901:2000].std() * lmrSTD + lmrMean
        tc_liz.pipe(func_norm)['normTCcounts'].plot(ax=ax, color='k', ls='-', lw=2, label='reconstructed HU')
    if modern_hurricane_on:
        lmrMean = ds_lmr[name].mean('mcrun').pipe(lowpass).sel(year=slice(1901,2000)).mean()
        lmrSTD = ds_lmr[name].mean('mcrun').pipe(lowpass).sel(year=slice(1901,2000)).std()
        func_norm = lambda x: (x - x.sel(year=slice(1901,2000)).mean())/x.sel(year=slice(1901,2000)).std() * lmrSTD + lmrMean
        ds_ah['AHlp'].pipe(func_norm).plot(ax=ax, color='gray', ls='--', lw=2, label='recorded HU')
        ds_ah['AMHlp'].pipe(func_norm).plot(ax=ax, color='k', ls='--', lw=2, label='recorded MH')
    if liz_add_on_lmr or modern_hurricane_on:
        ax.legend(frameon=False, loc='lower left', ncol=2)
        
    #ax.legend(loc='upper left', ncol=2)
    spines_off(ax)
    #ax.set_ylim(None, 7.5)
    ax.set_ylabel('HU #')
    ax.set_title(f'TC count from statistical model using {lmr} SST (lowpass_cutoff={lmr_window})', loc='left')
    
    #
    ax = axes[2]
    #MDR - TROP
    name = 'MDR$-$TROP'
    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=0.4)
    da.plot(ax=ax, label=name, lw=2)
    for r in ds_lmr.mcrun:
        da = (ds_lmr['MDR'] - ds_lmr['TROP']).sel(mcrun=r) \
            .pipe(lowpass) \
            .sel(year=slice(year0,None))
        da.plot(ax=ax, lw=1, color='C0', alpha=0.2)

    name = 'TROP'
    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=0.4)
    da.plot(ax=ax, label=name, lw=2)
    for r in ds_lmr.mcrun:
        da = ds_lmr[name].sel(mcrun=r) \
            .pipe(lowpass) \
            .sel(year=slice(year0,None))
        da.plot(ax=ax, lw=1, color='C1', alpha=0.2)
    
    name = 'MDR'
    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=0.4)
    da.plot(ax=ax, label=name, lw=2)
    for r in ds_lmr.mcrun:
        da = ds_lmr[name].sel(mcrun=r) \
            .pipe(lowpass) \
            .sel(year=slice(year0,None))
        da.plot(ax=ax, lw=1, color='C2', alpha=0.2)

    ax.legend(ncol=3, frameon=False)
    spines_off(ax)
    ax.set_title(f'MDR and tropical mean SSTs (lowpass_cutoff={lmr_window})', loc='left')
    ax.set_ylabel('K')
    ax.set_ylim(-0.4, 0.4)
    
    ax.set_xticks(range(900, 2000+1, 100));
    ax.set_xlabel('year')
    #ax.set_xlim(year0, 2015)
    ax.set_xlim(year0, 2000)
    
    #plt.tight_layout()
    figname = __file__.replace('.py', f'.png')
    if len(sys.argv)>1 and sys.argv[1]=='savefig':
        wysavefig(figname, dpi=200)
    else:
        print(f'figure not saved. to save figure: python {__file__} savefig')
    tt.check(f'**Done**')
    plt.show()