#!/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_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' #from gav
#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')
lmr = ifile.split('/')[-1].split('_')[0]
ds_lmr = xr.open_dataset(ifile).rename(MCrun='mcrun')

# fig, 1850
year0 = 850#851
year1 = 1850
n_window = 40#61
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))
#da = ds_lmr['HU'].filter.lowpass(1.0/n_window, dim='year', padtype='even').sel(year=slice(year0, year1)) #wy 2021-02-05
#MH instead of HU
mh = emulate_mh(ds_lmr, 'MDR', 'TROP', yearsRef=slice(1979,2000))
da = mh.filter.lowpass(1.0/n_window, dim='year', padtype='even').sel(year=slice(year0, year1)) #wy 2021-02-05

n_liz, n_mcrun = len(df.columns), da.mcrun.size
zz = np.zeros((n_liz, n_mcrun+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_mcrun):
        zz[ii, jj+1] = np.corrcoef(df.iloc[:, ii].values, da.isel(mcrun=jj).values)[0,1]
        n = ess_in_corr(df.iloc[:, ii].values, da.isel(mcrun=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('mcrun').values)[0,1]
    n = ess_in_corr(df.iloc[:, ii].values, da.mean('mcrun').values)
    zz_n[ii, 0] = n
    zz_thd[ii, 0] = p2r(alpha, n-2)

dazz = xr.DataArray(zz, dims=['jkmember', 'mcrun'], coords=[range(n_liz), range(n_mcrun+1)]).assign_attrs(long_name='correlation coefficient')
da_n = xr.DataArray(zz_n, dims=['jkmember', 'mcrun'], coords=[range(n_liz), range(n_mcrun+1)]).assign_attrs(long_name='effective sample size')
da_rthd = xr.DataArray(zz_thd, dims=['jkmember', 'mcrun'], coords=[range(n_liz), range(n_mcrun+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', 
              vmin=-1, vmax=1,
              yincrease=False)
    
    for ii in range(n_liz):
        for jj in range(n_mcrun+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_mcrun+1))
    ax.set_xticklabels(['em',]+list(range(1, n_mcrun+1)))
    ax.set_yticks(range(n_liz))
    ax.set_yticklabels(df.columns)
    ax.set_xlim(-0.5, n_mcrun+0.5)
    ax.set_ylim(n_liz-0.5, -0.5)
    ax.set_ylabel('')
    #ax.set_title(f'corr(jknife, {lmr}), {year0}-{year1}, lowpass cutoff period: {n_window} years')
    ax.set_title(f'corr(sedimentary reconstructions, LMR2.1 MHs), {year0}-{year1}')
    ax.set_xlabel('LMR ensemble member')

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    wyplot()
    
    """
    dazz.plot(figsize=figsize, 
              center=0, 
              #cmap='Spectral_r', 
              cmap='turbo', 
              vmin=-1, vmax=1,
              yincrease=False)
    
    for ii in range(n_liz):
        for jj in range(n_mcrun+1):
            if np.abs(zz[ii, jj]) >= zz_thd[ii, jj]:
                plt.text(jj, ii, f'{zz[ii, jj]*100:2.0f}', color='k',
                     ha='center', va='center')
            else:
                plt.text(jj, ii, f'{zz[ii, jj]*100:2.0f}', color='gray',
                     ha='center', va='center')
    plt.xticks(range(n_mcrun+1), ['em',]+list(range(1, n_mcrun+1)))
    plt.yticks(range(n_liz), df.columns)
    plt.xlim(-0.5, n_mcrun+0.5)
    plt.ylim(n_liz-0.5, -0.5)
    plt.ylabel('')
    plt.title(f'corr(jknife, {lmr}), {year0}-{year1}, lowpass cutoff period: {n_window} years')
    #plt.title(f'TC proxies with LMR TCs, {year0}-{year1}', loc='left')
    plt.xlabel('LMR ensemble member')
    """
    
    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()