#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Mar 11 15:51:41 EST 2021
#wy2021-05-28: ds_ah is obtained from data/humh_adjusted.nc
#wy2021-06-07: use the ax.secondary_yaxis method
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 fig_lines_tc_all import ds_liz, ds_jk, ds_ah, normalize, lowpass
#from fig_lines_tc_all import ds_liz, ds_jk, normalize, lowpass
from fig_lines_tc_all_poisson_HUcount import ds_liz, ds_jk, normalize, lowpass
from shared import spines_off
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'data/humh_adjusted.nc'
ds_ah = xr.open_dataset(ifile)
year0 = 850
alpha_fill = 0.2
#if True:
#    from wyconfig import *
#    ax = None
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:#target years in normalization
        norm_years = slice(1901,2000)
    ll, uu = ds_liz['smoothlower'].sel(year=slice(year0, None)), ds_liz['smoothupper'].sel(year=slice(year0, None))
    #normalize to have the same mean/std
    da = ds_liz['normTCcounts'].sel(year=slice(year0, None))
    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(ll_.year, ll_, uu_, alpha=alpha_fill, color='gray')
    #ds_liz['normTCcounts'].sel(year=slice(year0, None)).plot(ax = ax, color='gray', lw=2, label='sediment HU')
    da_.plot(ax = ax, color='gray', lw=2, label='sediment HU')
    #ax.axhline(ds_liz['normTCcounts'].sel(year=norm_years).mean(), color='gray', ls='--')
    #ax.axhline(da_.sel(year=norm_years).mean(), color='gray', ls='--')
    #jk members
    sites = ['noNewEng', 'noMidAtl', 'noSoutheast', 'noGulf_ChoBay', 'noGulf_AppBay',
        'noLIBHBahamas', 'noAbaco', 'noAndrosAM4', 'noVieqPR', 'noBelizeLRBH']
    for site in sites:
        #ds_jk[site].sel(year=slice(year0,None)).plot(ax=ax, lw=1, label=site, alpha=0.5)
        ds_jk[site].sel(year=slice(year0,None)) \
            .pipe(normalize, lowpass_mean=False, ensmean=da, yearsTarget=norm_years, multiplicative=norm_multiplicative) \
            .plot(ax=ax, lw=1, label=site, alpha=0.5)
    #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) \
    ci = da.pipe(lowpass) \
        .quantile([0.025, 0.975], dim='sample')
    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='--')


    #
    spines_off(ax)
    #ax.set_ylim(0, 0.07)
    #ax.set_ylim(2, 14)
    #ax.set_yticks(range(2,15,2))
    ax.set_ylim(3, 15)#wy2021-10-18
    ax.set_xlim(year0, 2000)
    ax.set_xticks(range(2000, 850, -100))
    ax.legend(ncol=4, frameon=False, loc='upper left')
    #ax.set_ylabel('normalized HU activity')
    ax.set_ylabel('HU count')

    """
    #right axis
    ax1 = 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')
    mhMean = da_mh.sel(year=slice(850,None)).pipe(lowpass).sel(year=norm_years).mean('year')
    mhSTD = da_mh.sel(year=slice(850,None)).pipe(lowpass).sel(year=norm_years).std('year')
    y0_right = (y0 - lizMean)/lizSTD*mhSTD + mhMean
    y1_right = (y1 - lizMean)/lizSTD*mhSTD + mhMean
    ax1.set_ylabel('modern MH count')#, color='C1')
    ax1.set_ylim(y0_right, y1_right)
    #ax1.spines['right'].set_linestyle('--')
    #ax1.legend(ncol=1, frameon=False, loc='upper right')
 
    #right axis, another one
    ax2 = 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
    ax2.set_ylabel('HU(MH) count')#, color='C1')
    ax2.set_ylim(y0_right, y1_right)
    yticks = ax2.get_yticks()
    yticklabels = [f'{ytick:2.0f}({(ytick - huMean)/huSTD*mhSTD + mhMean:.1f})' for ytick in yticks]
    ax2.set_yticklabels(yticklabels)
    #ax1.spines['right'].set_linestyle('--')
    #ax1.legend(ncol=1, frameon=False, loc='upper right')
    """
    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
    def hu2mh(y):
        return (y - huMean)/huSTD*mhSTD + mhMean
    def mh2hu(y):
        return (y - mhMean)/mhSTD*huSTD + huMean
    #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')
    ax_mh = ax.secondary_yaxis('right', functions=(hu2mh, mh2hu)) 
    ax_mh.set_ylabel('MH count')
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #norm_years = slice(1851,2000)#slice(1951,2000)#default is slice(1901,2000)
    #norm_years = slice(1901,2000)#slice(1951,2000)#default is slice(1901,2000)
    norm_years = slice(1870,2000)#wy2021-10-18: same as the one used in rMDR/TROP sst transform
    norm_multiplicative = False#True#default is False
    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)
    tt.check(f'**Done**')
    plt.show()