#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sun Jan 31 11:50:54 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
import matplotlib.pyplot as plt
#more imports
import xfilter
#from misc.cim import cim
from mystats import p2r
from shared import ess_in_corr, emulate_hu, emulate_mh
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
# params
# 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'
# 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'))
# lme tc
ifile = '/tigress/wenchang/analysis/LMTC/data/LME_TC.fullForcing.13ens.nc'
ds_lme = xr.open_dataset(ifile)
#WY2021-02-17: use years of reference 1982-2005 instead of 1901-2000
#ds_lme['HU'] = emulate_hu(ds_lme, 'MDR', 'TROP', yearsRef=slice(1982,2005), lmeCase=True)
ds_lme['MH'] = emulate_mh(ds_lme, 'MDR', 'TROP', yearsRef=slice(1982,2005), lmeCase=True)
#test
#import xfilter
#ds_lme['HU'].mean('en').filter.lowpass(1/40, dim='year', padtype='even').plot()
#plt.show()
#sys.exit()


# fig, 1850
year0 = 850#851
year1 = 1850
n_window = 40#61
padtype = 'even' # symmetric padding in lowpass filter
figsize = (8,3.3)

# lizzie
df = tc_jk
#sort the sites from north to south
sites = ['noNewEng', 'noMidAtl', 'noSoutheast', 'noGulf_ChoBay', 'noGulf_AppBay',
    'noLIBHBahamas', 'noAbaco', 'noAndrosAM4', 'noVieqPR', 'noBelizeLRBH']
df = df[sites]
df['normTCcounts'] = tc_liz['normTCcounts']
#df = df.sort_index(axis=1)
df = df[year0:year1]
df.head()

# lmr
#da = ds_lmr['HU'].filter.lowpass(1.0/n_window, dim='year').sel(year=slice(year0, year1))
# lme tc
#da = ds_lme['HU'].filter.lowpass(1.0/n_window, dim='year', padtype=padtype).sel(year=slice(year0, year1))
da = ds_lme['MH'].filter.lowpass(1.0/n_window, dim='year', padtype=padtype).sel(year=slice(year0, year1))

n_liz, n_en = len(df.columns), da.en.size
zz = np.zeros((n_liz, n_en+1)) # correlation coefficient between the two individual members from the two groups
zz_n = np.zeros_like(zz) # effective sample size
zz_thd = np.zeros_like(zz) # corr threshold value needed to be significant 
alpha = 0.05
for ii in range(n_liz):
    for jj in range(n_en):
        zz[ii, jj+1] = np.corrcoef(df.iloc[:, ii].values, da.isel(en=jj).values)[0,1]
        n = ess_in_corr(df.iloc[:, ii].values, da.isel(en=jj).values)
        zz_n[ii, jj+1] = n
        zz_thd[ii, jj+1] = p2r(alpha, n-2)
for ii in range(n_liz):
    zz[ii, 0] = np.corrcoef(df.iloc[:, ii].values, da.mean('en').values)[0,1]
    n = ess_in_corr(df.iloc[:, ii].values, da.mean('en').values)
    zz_n[ii, 0] = n
    zz_thd[ii, 0] = p2r(alpha, n-2)

dazz = xr.DataArray(zz, dims=['jkmember', 'en'], coords=[range(n_liz), range(n_en+1)]).assign_attrs(long_name='correlation coefficient')
da_n = xr.DataArray(zz_n, dims=['jkmember', 'en'], coords=[range(n_liz), range(n_en+1)]).assign_attrs(long_name='effective sample size')
da_rthd = xr.DataArray(zz_thd, dims=['jkmember', 'en'], coords=[range(n_liz), range(n_en+1)]).assign_attrs(long_name=f'corr. coef. threshold of significance at level {alpha}')
ds = xr.Dataset(dict(
    r=dazz,
    n_ess=da_n,
    r_thd=da_rthd
    ))

def wyplot(ax=None):
    if ax is None:
        fig,ax = plt.subplots(figsize=figsize)
    dazz.plot(ax=ax, 
              center=0, 
              #cmap='Spectral_r', 
              cmap='turbo', 
              #cmap='RdBu_r', 
              vmin=-1, vmax=1,
              yincrease=False)
    
    for ii in range(n_liz):
        for jj in range(n_en+1):
            if np.abs(zz[ii, jj]) >= zz_thd[ii, jj]:
                ax.text(jj, ii, f'{zz[ii, jj]*100:2.0f}', color='k',
                     ha='center', va='center')
            else:
                ax.text(jj, ii, f'{zz[ii, jj]*100:2.0f}', color='gray',
                     ha='center', va='center')
    ax.set_xticks(range(n_en+1))
    ax.set_xticklabels(['em',]+list(range(1, n_en+1)))
    ax.set_yticks(range(n_liz))
    ax.set_yticklabels(df.columns)
    ax.set_xlim(-0.5, n_en+0.5)
    ax.set_ylim(n_liz-0.5, -0.5)
    #ax.set_title(f'corr(jknife, LME), {year0}-{year1}, lowpass cutoff period: {n_window} years', loc='left')
    ax.set_title(f'corr(sedimentary reconstructions, LME MHs), {year0}-{year1}')
    ax.set_xlabel('LME ensemble member')
    ax.set_ylabel('')

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    wyplot()
    
    figname = __file__.replace('.py', f'.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        wysavefig(figname, dpi=128)
    else:
        print(f'figure not saved. to save figure: python {__file__} savefig')
    tt.check(f'**Done**')
    plt.show()