#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu May 27 11:55:29 EDT 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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
pastyears = slice(850,1849)
pdyears = slice(1901,2000)#slice(1850,2005)#slice(1961,2000)
futureyears = [ slice(2006,2100), slice(2061,2100) ][0]
key = 'MH'

#liz tc
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')))
da = ds_liz.normTCcounts
da_ref = da.sel(year=pastyears).mean('year')
daa_liz = ( da.sel(year=pdyears).mean('year') - da_ref )/da_ref * 100
#liz jk
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')))
da = ds_jk.to_array()
da_ref = da.sel(year=pastyears).mean('year')
daa_jk = ( da.sel(year=pdyears).mean('year') - da_ref )/da_ref * 100
#lmr2019
ifile = '/tigress/wenchang/analysis/LMTC/data/LMR2019_TC.nc'
ds = xr.open_dataset(ifile)
ds_lmr = ds.rename(MCrun='mcrun')
da = ds_lmr[key]
da_ref = da.sel(year=pastyears).mean('year')
daa_lmr2019 = ( da.sel(year=pdyears).mean('year') - da_ref )/da_ref * 100
#mcrun mean
da = da.mean('mcrun')
da_ref = da.sel(year=pastyears).mean('year')
daa_lmr2019_m = ( da.sel(year=pdyears).mean('year') - da_ref )/da_ref * 100
#lmr2018
ifile = '/tigress/wenchang/analysis/LMTC/data/LMR2018_TC.nc'
ds = xr.open_dataset(ifile)
ds_lmr = ds.rename(MCrun='mcrun')
da = ds_lmr[key]
da_ref = da.sel(year=pastyears).mean('year')
daa_lmr2018 = ( da.sel(year=pdyears).mean('year') - da_ref )/da_ref * 100
#mcrun
da = da.mean('mcrun')
da_ref = da.sel(year=pastyears).mean('year')
daa_lmr2018_m = ( da.sel(year=pdyears).mean('year') - da_ref )/da_ref * 100

#cmip5
ifile = 'ssttc.past1000histrcp85.0850-2100.ens06.nc'
ds = xr.open_dataset(ifile)
da = ds[key]
#% change from past to present day (PD)
refyears = pastyears
da_ref = da.sel(year=refyears).mean('year')
daa_pd = ( da.sel(year=pdyears).mean('year') - da_ref )/da_ref * 100 #% change from ref
daa_pd = daa_pd.assign_attrs(units='%', long_name=f'{key} # change from {refyears.start}-{refyears.stop} to {pdyears.start}-{pdyears.stop}')
#% change from PD to future
refyears = pdyears
da_ref = da.sel(year=refyears).mean('year')
daa_future = ( da.sel(year=futureyears).mean('year') - da_ref )/da_ref * 100 #% change from ref
daa_future = daa_future.assign_attrs(units='%', long_name=f'{key} # change from {refyears.start}-{refyears.stop} to {futureyears.start}-{futureyears.stop}')
ds = xr.Dataset({f'dN{key}pd':daa_pd, f'dN{key}future':daa_future})

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    constrained_layout_off()
    import xlinregress#,xaddon
    from shared import scaleax
    plt.close()
    #fig,ax = plt.subplots(figsize=(6,4.5), sharex=True)
    fig,axes = plt.subplots(2,1,figsize=(6,5), sharex=True)
    bins = np.arange(-20,11,2)
    alpha_hist = 0.8
    histtype = 'step'
    ax = axes[0]
    #ax.plot(daa_liz, 1, marker='*', ls='none', color='C0', label='sediment')
    ax.axvline(daa_liz, color='C0', ls='--')
    daa_jk.plot.hist(ax=ax, bins=bins, alpha=alpha_hist, color='C0', label='sediment', histtype=histtype)
    #ax.plot(daa_lmr2018_m, 1, marker='*', ls='none', color='C1', label='LMR2.0 ensm')
    ax.axvline(daa_lmr2018_m, color='C1', ls='--')
    daa_lmr2018.plot.hist(ax=ax, bins=bins, alpha=alpha_hist, color='C1', label='LMR2.0', histtype=histtype)
    #ax.plot(daa_lmr2019_m, 1, marker='*', ls='none', color='C2', label='LMR2.1 ensm')
    ax.axvline(daa_lmr2019_m, color='C2', ls='--')
    daa_lmr2019.plot.hist(ax=ax, bins=bins, alpha=alpha_hist, color='C2', label='LMR2.1', histtype=histtype)

    ax.legend(frameon=False)
    ax.set_xlabel('')
    ax.set_ylabel('hist #')
    ax.set_title('')
    ax.set_yticks(range(1,6,2))
    #ax.spines['top'].set_visible(False)
    #ax.spines['right'].set_visible(False)
    
    ax = axes[1]
    ds.plot.scatter(x=f'dN{key}pd', y=f'dN{key}future', ax=ax, color='none', edgecolors='k')
    for model in ds.model.values:
        #ax.text(ds[f'dN{key}pd'].sel(model=model), ds[f'dN{key}future'].sel(model=model), f'{model}\n', ha='center', va='center')
        ax.text(ds[f'dN{key}pd'].sel(model=model), ds[f'dN{key}future'].sel(model=model)-1, f'\n{model}', ha='center', va='top')
    ax.grid('on')
    ax.axhline(0, color='gray', ls='--')
    ax.axvline(0, color='gray', ls='--')
    ax.set_title('')
    #regression line
    rg = ds[f'dN{key}future'].linregress.numba(ds[f'dN{key}pd'])
    r,p = rg.r.item(), rg.pvalue.item()
    ax.plot(ds[f'dN{key}pd'], rg.intercept + rg.slope*ds[f'dN{key}pd'], alpha=0.5, color='gray')
    ax.text(20, 100, f'r = {r:.2f}\np = {p:.2f}', ha='right', va='center', color='gray')
    
    scaleax(axes[0], down=-0.9, up=0.3, left=-0.05)
    scaleax(axes[1], up=1, down=-0.05, left=-0.05)
    constrained_layout_on()
    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()