#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Jan 29 15:18:06 EST 2021
#wy2021-05-28: ds_ah is obtained from data/humh_adjusted.nc
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 shared import emulate_hu, emulate_mh, scaleax
#from data.data_adjusted_hurricanes import ds as ds_ah
#ds_ah = xr.open_dataset('data/humh_adjusted.nc')
#
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)
ds_liz = xr.Dataset(tc_liz).rename(years='year').pipe(lambda x: x.assign_coords(year=x.year.values.astype('int')))
# 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)
ds_jk = xr.Dataset(tc_jk).pipe(lambda x: x.assign_coords(year=x.year.values.astype('int')))

# tc lmr 2019
#ifile = '/tigress/gvecchi/DATA/LMR_2019/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', 'LMR2019_TC.nc')
ifile = '/tigress/wenchang/analysis/LMTC/data/LMR2019_TC.nc'
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
#MDR/TROP sst anomalies from years 1979-2000
func_anom = lambda x: x - x.sel(year=slice(1979,2000)).mean('year')
ds_lmr['MDR'] = ds_lmr['MDR'].pipe(func_anom)
ds_lmr['TROP'] = ds_lmr['TROP'].pipe(func_anom)
#WY2021-03-03: major hurricane(MH) use yearsRef of 1979-2000
mh = emulate_mh(ds_lmr, 'MDR', 'TROP', yearsRef=slice(1979,2000))
ds_lmr['MH'] = mh
#WY2021-03-03: major hurricane(MH) use yearsRef of 1979-2000
mh_noCO2 = emulate_mh(ds_lmr, 'MDR', 'TROP', yearsRef=slice(1979,2000), co2_on=False)
ds_lmr['MH_noCO2'] = mh_noCO2
"""

# 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')
ifile = '/tigress/wenchang/analysis/LMTC/data/LMR2018_TC.nc'
ds = xr.open_dataset(ifile)
ds_lmr2018 = ds.rename(MCrun='mcrun')
"""
#WY2021-02-17: use yearsRef of 1979-2000 (equivalent to 1982-2005 for HadISST)
hu = emulate_hu(ds_lmr2018, 'MDR', 'TROP', yearsRef=slice(1979,2000))
ds_lmr2018['HU'] = hu
#MDR/TROP sst anomalies from years 1979-2000
func_anom = lambda x: x - x.sel(year=slice(1979,2000)).mean('year')
ds_lmr2018['MDR'] = ds_lmr2018['MDR'].pipe(func_anom)
ds_lmr2018['TROP'] = ds_lmr2018['TROP'].pipe(func_anom)
#WY2021-03-03: major hurricane(MH) use yearsRef of 1979-2000
mh = emulate_mh(ds_lmr2018, 'MDR', 'TROP', yearsRef=slice(1979,2000))
ds_lmr2018['MH'] = mh
#WY2021-03-03: major hurricane(MH) use yearsRef of 1979-2000
mh_noCO2 = emulate_mh(ds_lmr2018, 'MDR', 'TROP', yearsRef=slice(1979,2000), co2_on=False)
ds_lmr2018['MH_noCO2'] = mh_noCO2
"""

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 normalize(da, lowpass_on=True, ensmean=None, yearsTarget=None, multiplicative=False):
    """normalize so that the input dataset/dataarray has the same mean and std as the target"""
    if lowpass_on:
        da = da.pipe(lowpass)
    daname = 'normTCcounts'
    if yearsTarget is None:
        yearsTarget = slice(1901,2000)
    targetMean = ds_liz[daname].sel(year=yearsTarget).mean('year')
    targetSTD = ds_liz[daname].sel(year=yearsTarget).std('year')
    if ensmean is None:
        da_ensmean = da
    else:
        da_ensmean = ensmean
        if lowpass_on:
            da_ensmean = lowpass(da_ensmean)
    if multiplicative:
        return da/da_ensmean.sel(year=yearsTarget).mean('year') * targetMean
    return (da - da_ensmean.sel(year=yearsTarget).mean('year'))/da_ensmean.sel(year=yearsTarget).std('year') * targetSTD + targetMean

def wyplot(ax=None, norm_years=None, norm_multiplicative=False, LMRMHon=True): 
    if ax is None:
        fig, ax = plt.subplots(figsize=(8.5,3))
    if norm_years is None:#years for normalization
        norm_years = slice(1901,2000)
    #Liz
    da, ll, uu = ds_liz['normTCcounts'], ds_liz['smoothlower'], ds_liz['smoothupper']
    ax.fill_between(da.year, ll, uu, alpha=alpha_fill, color='gray')
    da.plot(ax=ax, label='sediment HU', lw=2, color='gray')
    #modern HU
    da = ds_ah['HU']
    da_hu = da.mean('sample')
    ci = da.pipe(normalize, ensmean=da.mean('sample'), yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        .quantile([0.025, 0.975], dim='sample')
    ax.fill_between(ci.year, ci[0].values, ci[1].values, alpha=alpha_fill/2, color='k')
    """
    stde = da.pipe(normalize, ensmean=da.mean('sample'), yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        .std('sample')
    da_ = da.mean('sample').pipe(normalize, yearsTarget=norm_years, multiplicative=norm_multiplicative)
    ci = da_ - stde, da_ + stde
    ax.fill_between(da_.year, ci[0].values, ci[1].values, alpha=alpha_fill/2, color='k')
    """
    da.mean('sample').pipe(normalize, yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        .plot(ax=ax, label='modern HU', lw=2, color='k')
    #modern MH
    da = ds_ah['MH']
    da_mh = da.mean('sample')
    ci = da.pipe(normalize, ensmean=da.mean('sample'), yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        .quantile([0.025, 0.975], dim='sample')
    ax.fill_between(ci.year, ci[0], ci[1], alpha=alpha_fill/2, color='k')
    """
    stde = da.pipe(normalize, ensmean=da.mean('sample'), yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        .std('sample')
    da_ = da.mean('sample').pipe(normalize, yearsTarget=norm_years, multiplicative=norm_multiplicative)
    ci = da_ - stde, da_ + stde
    ax.fill_between(da_.year, ci[0], ci[1], alpha=alpha_fill/2, color='k')
    """
    da.mean('sample').pipe(normalize, yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        .plot(ax=ax, label='modern MH', lw=2, color='k', ls='--')

    #HU LMR2018
    name = 'HU'
    ds = ds_lmr2018
    da = ds[name].mean('mcrun') \
        .pipe(normalize, yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        .sel(year=slice(year0,None))
    #monte carlo simulation using Poisson distribution
    rng = np.random.default_rng(0)
    qq = [0.25, 0.75]
    da_ = ds[name].sel(year=slice(year0,None))
    zz = rng.poisson(da_, size=(10000,)+da_.shape)
    da_ = xr.DataArray(zz, dims=('sample',)+da_.dims).assign_coords(year=da_.year.values, mcrun=da_.mcrun.values).stack(s=('sample', 'mcrun'))
    #ci = da_.pipe(normalize, ensmean=da_.mean('s'), yearsTarget=norm_years, multiplicative=norm_multiplicative) \
    #    .quantile(qq, dim='s')
    stde = da_.pipe(normalize, ensmean=da_.mean('s'), yearsTarget=norm_years, multiplicative=norm_multiplicative).std('s')
    ci = da - stde, da + stde
    """
    da_ = xr.DataArray(zz, dims=('sample',)+da_.dims).assign_coords(year=da_.year.values, mcrun=da_.mcrun.values).mean('mcrun')
    ci = da_.pipe(normalize, ensmean=da_.mean('sample'), yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        .quantile([0.025, 0.975], dim='sample')
    """
    ax.fill_between(da.year, ci[0], ci[1], alpha=alpha_fill, color='C0')#, label='%95 CI')
    da.plot(ax=ax, label='LMR2.0SST-emulated HU', lw=2, color='C0')
    ax.axhline(da.sel(year=norm_years).mean(), color='gray', ls='--')
    if LMRMHon:
        #MH LMR2018
        name = 'MH'
        ds = ds_lmr2018
        da = ds[name].mean('mcrun') \
            .sel(year=slice(year0,None)) \
            .pipe(normalize, yearsTarget=norm_years, multiplicative=norm_multiplicative) 
        #monte carlo simulation using Poisson distribution
        #rng = np.random.default_rng(0)
        da_ = ds[name].sel(year=slice(year0,None))
        zz = rng.poisson(da_, size=(10000,)+da_.shape)
        da_ = xr.DataArray(zz, dims=('sample',)+da_.dims).assign_coords(year=da_.year.values, mcrun=da_.mcrun.values).stack(s=('sample','mcrun'))
        #ci = da_.pipe(normalize, ensmean=da_.mean('s'), yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        #    .quantile(qq, dim='s')
        stde = da_.pipe(normalize, ensmean=da_.mean('s'), yearsTarget=norm_years, multiplicative=norm_multiplicative).std('s')
        ci = da - stde, da + stde
        """
        da_ = xr.DataArray(zz, dims=('sample',)+da_.dims).assign_coords(year=da_.year.values, mcrun=da_.mcrun.values).mean('mcrun')
        ci = da_.pipe(normalize, ensmean=da_.mean('sample'), yearsTarget=norm_years, multiplicative=norm_multiplicative) \
            .quantile([0.025, 0.975], dim='sample')
        """
        ax.fill_between(da.year, ci[0], ci[1], alpha=alpha_fill, color='C0')#, label='%95 CI')
        da.plot(ax=ax, label='LMR2.0SST-emulated MH', lw=2, color='C0', ls='--')

    #HU LMR2019
    name = 'HU'
    ds = ds_lmr
    da = ds[name].mean('mcrun') \
        .pipe(normalize, yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        .sel(year=slice(year0,None))
    #monte carlo simulation using Poisson distribution
    #rng = np.random.default_rng(0)
    da_ = ds[name].sel(year=slice(year0,None))
    zz = rng.poisson(da_, size=(10000,)+da_.shape)
    da_ = xr.DataArray(zz, dims=('sample',)+da_.dims).assign_coords(year=da_.year.values, mcrun=da_.mcrun.values).stack(s=('sample','mcrun'))
    #ci = da_.pipe(normalize, ensmean=da_.mean('s'), yearsTarget=norm_years, multiplicative=norm_multiplicative) \
    #    .quantile(qq, dim='s')
    stde = da_.pipe(normalize, ensmean=da_.mean('s'), yearsTarget=norm_years, multiplicative=norm_multiplicative).std('s')
    ci = da - stde, da + stde
    """
    da_ = xr.DataArray(zz, dims=('sample',)+da_.dims).assign_coords(year=da_.year.values, mcrun=da_.mcrun.values).mean('mcrun')
    ci = da_.pipe(normalize, ensmean=da_.mean('sample'), yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        .quantile([0.025, 0.975], dim='sample')
    """
    ax.fill_between(da.year, ci[0], ci[1], alpha=alpha_fill, color='C1')#, label='%95 CI')
    da.plot(ax=ax, label='LMR2.1SST-emulated HU', lw=2, color='C1')
    #MH LMR2019
    if LMRMHon:
        name = 'MH'
        ds = ds_lmr
        da = ds[name].mean('mcrun') \
            .sel(year=slice(year0,None)) \
            .pipe(normalize, yearsTarget=norm_years, multiplicative=norm_multiplicative) 
        #monte carlo simulation using Poisson distribution
        #rng = np.random.default_rng(0)
        da_ = ds[name].sel(year=slice(year0,None))
        zz = rng.poisson(da_, size=(10000,)+da_.shape)
        da_ = xr.DataArray(zz, dims=('sample',)+da_.dims).assign_coords(year=da_.year.values, mcrun=da_.mcrun.values).stack(s=('sample','mcrun'))
        #ci = da_.pipe(normalize, ensmean=da_.mean('s'), yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        #    .quantile(qq, dim='s')
        stde = da_.pipe(normalize, ensmean=da_.mean('s'), yearsTarget=norm_years, multiplicative=norm_multiplicative).std('s')
        ci = da - stde, da + stde
        """
        da_ = xr.DataArray(zz, dims=('sample',)+da_.dims).assign_coords(year=da_.year.values, mcrun=da_.mcrun.values).mean('mcrun')
        ci = da_.pipe(normalize, ensmean=da_.mean('sample'), yearsTarget=norm_years, multiplicative=norm_multiplicative) \
            .quantile([0.025, 0.975], dim='sample')
        """
        ax.fill_between(da.year, ci[0], ci[1], alpha=alpha_fill, color='C1')#, label='%95 CI')
        da.plot(ax=ax, label='LMR2.0SST-emulated MH', lw=2, color='C1', ls='--')

    #ax.legend(loc='upper left', ncol=2)
    spines_off(ax)
    ax.set_ylim(0, 0.07)
    ax.set_ylabel('normalized HU activity')
    #ax.set_title(f'TC count from statistical model using LMR SST (lowpass_cutoff={lmr_window})', loc='left')
    
    #
    ax.set_xticks(range(900, 2000+1, 100));
    ax.set_xlabel('year')
    #ax.set_xlim(year0, 2015)
    ax.set_xlim(year0, 2000)
    ax.legend(ncol=3, frameon=False, loc='upper right')
    
    """
    #right axis
    ax1 = ax.twinx()
    y0, y1 = ax.get_ylim()
    lizMean = ds_liz.normTCcounts.sel(year=norm_years).mean('year').item()
    lizSTD = ds_liz.normTCcounts.sel(year=norm_years).std('year').item()
    huMean = da_hu.sel(year=slice(850,None)).pipe(lowpass).sel(year=norm_years).mean('year').item()
    huSTD = da_hu.sel(year=slice(850,None)).pipe(lowpass).sel(year=norm_years).std('year').item()
    mhMean = da_mh.sel(year=slice(850,None)).pipe(lowpass).sel(year=norm_years).mean('year').item()
    mhSTD = da_mh.sel(year=slice(850,None)).pipe(lowpass).sel(year=norm_years).std('year').item()
    y0_right = (y0 - lizMean)/lizSTD*huSTD + huMean
    y1_right = (y1 - lizMean)/lizSTD*huSTD + huMean
    ax1.set_ylim(y0_right, y1_right)
    ax1.set_ylabel('HU(MH) count')#, color='C1')
    yticks = ax1.get_yticks()
    yticklabels = [f'{ytick:2.0f}({(ytick - huMean)/huSTD*mhSTD + mhMean:.1f})' for ytick in yticks]
    ax1.set_yticklabels(yticklabels)
    #ax1.spines['right'].set_linestyle('--')

    #right axis, another one
    #x,y,w,h = ax.get_position().bounds
    #ax2 = plt.gcf().add_axes([x+w, y, 0.1, h]) 
    #ax2.yaxis.tick_right()
    #scaleax(ax2, right=0.2, left=0.2)
    #ax2.spines['left'].set_visible(False)
    #ax2.spines['top'].set_visible(False)
    #ax2.spines['bottom'].set_visible(False)
    ax2 = ax.twinx()
    y0, y1 = ax.get_ylim()
    lizMean = ds_liz.normTCcounts.sel(year=norm_years).mean('year')
    lizSTD = ds_liz.normTCcounts.sel(year=norm_years).std('year')
    huMean = da_hu.sel(year=slice(850,None)).pipe(lowpass).sel(year=norm_years).mean('year')
    huSTD = da_hu.sel(year=slice(850,None)).pipe(lowpass).sel(year=norm_years).std('year')
    y0_right = (y0 - lizMean)/lizSTD*huSTD + huMean
    y1_right = (y1 - lizMean)/lizSTD*huSTD + huMean
    ax2.set_ylim(y0_right, y1_right)
    ax2.set_ylabel('HU count')#, color='C1')
    #ax2.spines['right'].set_linestyle('--')
    #ax2.set_yticks(np.arange(5,11.1,2))
    """
    lizMean = ds_liz.normTCcounts.sel(year=norm_years).mean('year').item()
    lizSTD = ds_liz.normTCcounts.sel(year=norm_years).std('year').item()
    huMean = da_hu.sel(year=slice(850,None)).pipe(lowpass).sel(year=norm_years).mean('year').item()
    huSTD = da_hu.sel(year=slice(850,None)).pipe(lowpass).sel(year=norm_years).std('year').item()
    mhMean = da_mh.sel(year=slice(850,None)).pipe(lowpass).sel(year=norm_years).mean('year').item()
    mhSTD = da_mh.sel(year=slice(850,None)).pipe(lowpass).sel(year=norm_years).std('year').item()
    def sediment2hu(y):
        return (y - lizMean)/lizSTD*huSTD + huMean
    def hu2sediment(y):
        return (y - huMean)/huSTD*lizSTD + lizMean
    def sediment2mh(y):
        return (y - lizMean)/lizSTD*mhSTD + mhMean
    def mh2sediment(y):
        return (y - mhMean)/mhSTD*lizSTD + lizMean
    ax_hu = ax.secondary_yaxis('right', functions=(sediment2hu, hu2sediment))
    ax_hu.set_ylabel('HU count')
    ax_mh = ax.secondary_yaxis(1.1, functions=(sediment2mh, mh2sediment))
    ax_mh.set_ylabel('MH count')

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    # fig
    norm_years = slice(1851,2000)
    norm_multiplicative = False#True
    LMRMHon = False
    wyplot(norm_years=norm_years, norm_multiplicative=norm_multiplicative, LMRMHon=LMRMHon)

    figname = __file__.replace('.py', f'.png')
    if norm_multiplicative:
        figname = figname.replace('.png', f'_multiplicative{norm_years.start}-{norm_years.stop}.png')
    else:
        figname = figname.replace('.png', f'_norm{norm_years.start}-{norm_years.stop}.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()