#!/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')
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')
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

# tc from transformed MDR/TROP of lmr2019
ifile = os.path.join(os.path.dirname(__file__), 'data', 'LMR_transformed', 'LMR2019_ssttc_transformed.nc')
ds = xr.open_dataset(ifile)
ds_lmr2019tf = ds.rename(MCrun='mcrun')

# tc from transformed MDR/TROP of lmr2018
ifile = os.path.join(os.path.dirname(__file__), 'data', 'LMR_transformed', 'LMR2018_ssttc_transformed.nc')
ds = xr.open_dataset(ifile)
ds_lmr2018tf = ds.rename(MCrun='mcrun')

year0 = 850
lmr_window = 40#61
alpha_fill = 0.1
# 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, ensmean=None, yearsTarget=None, multiplicative=False, lowpass_mean=True):
    """normalize so that the input dataset/dataarray has the same mean and std as the target"""
    if yearsTarget is None:
        yearsTarget = slice(1901,2000)
    """
    daname = 'normTCcounts'
    targetMean = ds_liz[daname].sel(year=yearsTarget).mean('year')
    targetSTD = ds_liz[daname].sel(year=yearsTarget).std('year')
    """
    #target to modern HU
    targetMean = ds_ah['HU'].mean('sample').pipe(lowpass).sel(year=yearsTarget).mean('year')
    targetSTD = ds_ah['HU'].mean('sample').pipe(lowpass).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 lowpass_mean:
        selfMean = da_ensmean.pipe(lowpass).sel(year=yearsTarget).mean('year')
        selfSTD = da_ensmean.pipe(lowpass).sel(year=yearsTarget).std('year')
    else:
        selfMean = da_ensmean.sel(year=yearsTarget).mean('year')
        selfSTD = da_ensmean.sel(year=yearsTarget).std('year')
    #print(f'{selfMean = }; {selfSTD = }; {targetMean = }; {targetSTD = }')
    if multiplicative:
        return da/selfMean * targetMean
    return (da - selfMean)/selfSTD * targetSTD + targetMean

def wyplot(ax=None, norm_years=None, norm_multiplicative=False): 
    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']
    #normalize to have the same mean/std 
    da_ = da.pipe(normalize, lowpass_mean=False, yearsTarget=norm_years, multiplicative=norm_multiplicative)
    ll_ = ll.pipe(normalize, lowpass_mean=False, ensmean=da, yearsTarget=norm_years, multiplicative=norm_multiplicative)
    uu_ = uu.pipe(normalize, lowpass_mean=False, ensmean=da, yearsTarget=norm_years, multiplicative=norm_multiplicative)
    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')
    ci = da.pipe(lowpass).quantile([0.025, 0.975], dim='sample') #lowpass instead of normalize
    ax.fill_between(ci.year, ci[0].values, ci[1].values, alpha=alpha_fill/2, color='k')
    #da.mean('sample').pipe(normalize, yearsTarget=norm_years, multiplicative=norm_multiplicative) \
    da.mean('sample').pipe(lowpass) \
        .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) \
        .pipe(lowpass) \
        .quantile([0.025, 0.975], dim='sample')
    ax.fill_between(ci.year, ci[0], ci[1], alpha=alpha_fill/2, color='k')
    da.mean('sample').pipe(normalize, yearsTarget=norm_years, multiplicative=norm_multiplicative) \
        .pipe(lowpass) \
        .plot(ax=ax, label='modern MH', lw=2, color='k', ls='--')

    #HU LMR2018
    name = 'HU'
    print(name, 2018)
    ds = ds_lmr2018tf
    da = ds[name].mean('mcrun') \
        .pipe(lowpass) \
        .sel(year=slice(year0,None))
    #monte carlo simulation using Poisson distribution
    rng = np.random.default_rng(0)
    qq = [0.025, 0.975]
    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(lowpass).quantile(qq, dim='s')
    ax.fill_between(ci.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='--')

    #HU LMR2019
    name = 'HU'
    print(name, 2019)
    ds = ds_lmr2019tf
    da = ds[name].mean('mcrun') \
        .pipe(lowpass) \
        .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(lowpass).quantile(qq, dim='s')
    ax.fill_between(ci.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')

    #ax.legend(loc='upper left', ncol=2)
    spines_off(ax)
    #ax.set_ylim(None, 0.07)
    #ax.set_ylim(2, 14)
    ax.set_ylim(4, 10)
    #ax.set_yticks(range(2,15,2))
    #ax.set_ylabel('normalized HU count')
    ax.set_ylabel('HU count')
    #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 - huMean)/huSTD*mhSTD + mhMean
    y1_right = (y1 - huMean)/huSTD*mhSTD + mhMean
    ax1.set_ylim(y0_right, y1_right)
    ax1.set_ylabel('MH count')#, color='C1')
    yticks = ax1.get_yticks()
    yticklabels = [f'{ytick:2.0f}' for ytick in yticks]
    ax1.set_yticklabels(yticklabels)
    ax1.spines['right'].set_visible(True)
    ax1.grid(False)


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

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