'''
This script is used to analyze PDP and SHAP modeling results
'''

import os,glob,cmaps,warnings,sys,pickle,re
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
import hydroeval as he
import matplotlib.ticker as ticker
import scienceplots
from statsmodels.nonparametric.kernel_regression import KernelReg
import shap
from matplotlib.lines import Line2D
import fasttreeshap
from src.myshap import createShapExplanation
from sklearn.ensemble import RandomForestRegressor
import multiprocessing as mp
import scipy as sp
import seaborn as sns
from matplotlib.colors import LogNorm, Normalize
warnings.filterwarnings("ignore")

# set custom font (not a good idea)
import matplotlib.font_manager as font_manager
font_manager.fontManager.addfont(os.environ['DATA']+'/fonts/Helvetica/Helvetica.ttf')
font_manager.fontManager.addfont(os.environ['DATA']+'/fonts/Helvetica/Helvetica-Bold.ttf')

plt.style.use(['science','nature','no-latex']) # require install SciencePlots
plt.rcParams.update({"font.size":12, 'font.family':'Helvetica'}) 

par_name = pd.read_csv('../data/predictors.csv')

climateVariable = sys.argv[1]
if climateVariable not in ['AnnSum','Pvol','Rx1day','Rx5day']:
    sys.exit('climateVariable must be one of [\'AnnSum\',\'Pvol\',\'Rx1day\',\'Rx5day\']')

############################################################################################
# read climate observation
fnames = glob.glob('../results/*/cv_scatter*_log_500km2_addLonLat*.csv')
df_sim = pd.concat([pd.read_csv(a).assign(model=a.split('/')[2]) for a in fnames])
reach_ids = pd.read_csv(fnames[0]).reach_id.values
df_sim[['tasmax','tasmin']] = df_sim[['tasmax', 'tasmin']] - 273.15

# read climate observation
fnames = glob.glob('../data/GRIT_catch_mean_1981-2014_obs_climate_MSWEP_CRU_TerraClimate_*.csv')
def func(fname):
    df0 = pd.read_csv(fname, index_col = [0]).set_index(['year','name']).T
    df0 = df0.loc[df0.index.astype(int).isin(reach_ids),:]
    # calculate mean of climate observation in 1981-2014
    df0 = df0.T.reset_index()
    df0 = df0.loc[(df0.year>=1981)&(df0.year<=2014),:].drop(columns=['year'])
    df0 = df0.groupby('name').mean().T.reset_index().rename(columns={'index':'reach_id'})
    df0['reach_id'] = df0.reach_id.astype(int)
    print(fname)
    return(df0)
pool = mp.Pool(16)
df_obs = pool.map(func, fnames)
df_obs = pd.concat(df_obs)

df_obs = df_obs.rename(columns={'tmn':'tasmin',
                                'tmx':'tasmax',
                                'ws':'sfcWind',
                                'srad':'rsds'})

# merge df_sim and df_obs
df_all = df_sim.merge(df_obs[['reach_id', 'AnnSum', 'Pvol', 'Rx1day', 'Rx5day', 'tasmin', 'tasmax']], on = 'reach_id', suffixes = ['_sim','_obs'])

# calculate bias of precipitation indices
for i in ['AnnSum', 'Pvol', 'Rx1day', 'Rx5day', 'tasmin', 'tasmax']:
    df_all[i+'_abs_bias'] = df_all[i+'_sim'] - df_all[i+'_obs']
    df_all[i+'_rel_bias'] = df_all[i+'_abs_bias'] / df_all[i+'_obs'] * 100

############################################################################################
# RF modeling and SHAP testing for a single GCM
def shap_test(model):
    # read SHAP
    fname = glob.glob('../results/*/shap_object*%s*p[0-9]f[0-9].pkl'%model)[0]
    with open(fname, 'rb') as file:  
        shap_values1 = pickle.load(file)
    # subset df_all
    df = df_all.loc[df_all.model==model,:]
    ens = os.path.basename(glob.glob('../results/%s/cv_scatter*addLonLat*csv'%model)[0]).split('_')[3]
    # set predictand of RF modeling
    y = df['RL10'].astype(float).values
    y = y / df['ohdb_catchment_area'].values
    y = np.log(y)
    df['target'] = y
    # set predictors of RF modeling
    climate_predictors = ['AnnSum', 'Pvol', 'Rx1day', 'Rx5day', 'hurs', 'rsds', 'sfcWind', 'tasmax', 'tasmin', 'snw',
                            'AnnSum_std', 'Pvol_std', 'Rx1day_std', 'Rx5day_std']
    geo_predictors = [  'ai', 'elevation', 'et0', 'arid', 'cold', 'temperate',
                        'tropical', 'forest', 'grassland', 'urban', 'slope', 'soilDep', 'GLHYMPS', 'lon', 'lat']
    other_predictors = ['FI', 'stream_elas', 'BFI', 'runoff_ratio', 'Q10_50', 'high_q_freq', 'cv']
    feature_names1 = [a+'_sim' if a not in df.columns else a for a in climate_predictors] + geo_predictors + other_predictors  
    map0 = lambda x: par_name.loc[par_name.par==(x[:-4] if '_sim' in x or '_obs' in x else x), 'name'].values[0]
    X1 = df[feature_names1].rename(columns=map0)
    # set predictors of RF modeling with the 1st climate predictors replaced
    feature_names2 = [re.sub('_sim','_obs',a) if climateVariable in a else a for a in feature_names1]  # replace the climateVariable by observation
    # feature_names2 = [re.sub('_sim','_obs',a) for a in feature_names1]  # replace the climateVariable by all observations
    X2 = df[feature_names2].rename(columns=map0)
    # run RF modelling by using climate observations or simulations
    clf = RandomForestRegressor(n_jobs = 24, n_estimators = 500)
    for name,x0 in zip(['_raw','_new'], [X1, X2]):
        clf.fit(x0, y)
        df['pred'+name] = clf.predict(x0)
        df['bias_rel'+name] = (df['pred'+name] - df['target']) / df['target'] * 100
        df['bias_abs'+name] = (df['pred'+name] - df['target'])
    # SHAP analysis
    if not os.path.exists('../results/%s/shap_object_%s_%s_obs_%s.pkl'%(model,model,ens,climateVariable)):
        explainer = fasttreeshap.TreeExplainer(clf, algorithm = 'auto', n_jobs = -1)
        shap_values0 = explainer(X2)
        shap_values0 = createShapExplanation(shap_values0)
        with open('../results/%s/shap_object_%s_%s_obs_%s.pkl'%(model,model,ens,climateVariable),'wb') as f:
            pickle.dump(shap_values0, f)
        with open('../results/%s/shap_explainer_object_%s_%s_obs_%s.pkl'%(model,model,ens,climateVariable), 'wb') as f:  
            pickle.dump(explainer, f)
    else:
        with open('../results/%s/shap_object_%s_%s_obs_%s.pkl'%(model,model,ens,climateVariable),'rb') as f:
            shap_values0 = pickle.load(f)
    shap_values2 = shap_values0

    # insert original and new SHAP values
    df = df.set_index('reach_id').loc[reach_ids].reset_index()
    for name in shap_values1.feature_names:
        df[name.strip()+'_shap_raw'] = shap_values1[:,name].values
        df[name.strip()+'_shap_new'] = shap_values2[:,name].values
        df[name.strip()+'_shap_delta'] = df[name.strip()+'_shap_raw'] - df[name.strip()+'_shap_new']
    
    return df

# lowess modeling
def lowess(x, y, f=1./3.):
    """
    Basic LOWESS smoother with uncertainty. 
    Note:
        - Not robust (so no iteration) and
             only normally distributed errors. 
        - No higher order polynomials d=1 
            so linear smoother.
    """
    # get some paras
    xwidth = f*(x.max()-x.min()) # effective width after reduction factor
    N = len(x) # number of obs
    # Don't assume the data is sorted
    order = np.argsort(x)
    # storage
    y_sm = np.zeros_like(y)
    y_stderr = np.zeros_like(y)
    # define the weigthing function -- clipping too!
    tricube = lambda d : np.clip((1- np.abs(d)**3)**3, 0, 1)
    # run the regression for each observation i
    for i in range(N):
        dist = np.abs((x[order][i]-x[order]))/xwidth
        w = tricube(dist)
        # form linear system with the weights
        A = np.stack([w, x[order]*w]).T
        b = w * y[order]
        ATA = A.T.dot(A)
        ATb = A.T.dot(b)
        # solve the syste
        sol = np.linalg.solve(ATA, ATb)
        # predict for the observation only
        yest = A[i].dot(sol)# equiv of A.dot(yest) just for k
        place = order[i]
        y_sm[place]=yest 
    x_sm = x[order]
    y_sm = y_sm[order]
    return x_sm, y_sm

# plot function
def plot_pdp(ax, df, varX, legend = True, ms = 2, fontsize = 8):

    varXname = par_name.loc[par_name.par==varX,'name'].values[0].strip()

    # SHAP-based PDP plots
    colNames1 = [varXname + '_shap_raw'] + [varX+'_sim' if varX+'_sim' in df.columns else varX]
    colNames2 = [varXname + '_shap_new'] + [varX+'_obs' if varX+'_sim' in df.columns else varX]

    df1 = df[colNames1].assign(type='CMIP6 only')
    df2 = df[colNames2].assign(type='CMIP6 + observed ' + climateVariable)
    df1 = df1.rename(columns={varX+'_sim':varXname} | {varX:varXname, varXname + '_shap_raw':varXname + '_shap'})
    df2 = df2.rename(columns={varX+'_obs':varXname} | {varX:varXname, varXname + '_shap_new':varXname + '_shap'})

    p1 = ax.scatter(df1[varXname].values, df1[varXname+'_shap'].values, s = ms, label = 'CMIP6 only', alpha = .2, color = '#1bd433')
    p2 = ax.scatter(df2[varXname].values, df2[varXname+'_shap'].values, s = ms, label = 'CMIP6 + observed ' + climateVariable, alpha = .2, color = '#4a9efe')
    
    ax.set_xlabel(varXname, fontsize = fontsize)
    ax.set_ylabel('SHAP', fontsize = fontsize)

    # LOWESS regression
    x_sm, y_sm = lowess(df1[varXname].values, df1[varXname+'_shap'].values, f=1./3.)
    ax.plot(x_sm, y_sm, c = '#2e9f5b', lw = 2) # p1.get_facecolors()
    x_sm, y_sm = lowess(df2[varXname].values, df2[varXname+'_shap'].values, f=1./3.)
    ax.plot(x_sm, y_sm, c = '#4157de', lw = 2)

    if legend:
        line1 = Line2D([], [], color='#2e9f5b', ls="-", linewidth=1.5)
        line2 = Line2D([], [], color='#4157de', ls="-", linewidth=1.5)
        sc1 = plt.scatter([],[], s=15, facecolors=p1.get_ec(), edgecolors=p1.get_ec())
        sc2 = plt.scatter([],[], s=15, facecolors=p2.get_ec(), edgecolors=p2.get_ec())
        ax.legend([(sc1,line1), (sc2,line2)], ['CMIP6 only','CMIP6 + observed ' + climateVariable], numpoints=1, handlelength = 1)

def scatter_shap_bias(ax, df, varBias, varDeltaSHAPs, varHue, ms = 2, fontsize = 8):

    if isinstance(varDeltaSHAPs, str):
        varName = par_name.loc[par_name.par==varDeltaSHAPs,'name'].values[0]
        ylabel = 'ΔSHAP for ' + varName.strip().lower()
    if isinstance(varDeltaSHAPs, int):
        ylabel = 'Sum of ΔSHAP for all basin attributes\nin top-%d important features'%varDeltaSHAPs

    # ΔSHAP for all basin attributes vs Rx1day bias
    if isinstance(varDeltaSHAPs, int):
        tmp = df.loc[:,df.columns.str.contains('_shap_raw')].abs().mean().sort_values(ascending = False) # get variable importance
        tmp = [a.split('_')[0] for a in tmp.iloc[:varDeltaSHAPs].index.values] # focus on top-5 important variables
        df['tmp'] = df[[a+'_shap_delta' for a in tmp if a in par_name.loc[par_name.type!='CMIP6','name'].values]].sum(1)
        sns.regplot(data = df, x = varBias+'_abs_bias', y = 'tmp', ax = ax, line_kws=dict(color="r"), scatter = False)
        if varHue == 'climate':
            sns.scatterplot(data = df, x = varBias+'_abs_bias', y = 'tmp', ax = ax, hue = varHue, s = ms, palette = ["#65C2A5",'#C7B18A',"#A89BAF",'#F8D347'])
        else:
            sns.scatterplot(data = df, x = varBias+'_abs_bias', y = 'tmp', ax = ax, hue = varHue, s = ms)
        r, p = sp.stats.pearsonr(df[varBias+'_abs_bias'].values, df['tmp'].values) 
    else:
        # ΔSHAP vs single-variable bias
        sns.regplot(data = df, x = varBias + '_abs_bias', y = varName + '_shap_delta', ax = ax, line_kws=dict(color="r"), scatter = False)
        if varHue == 'climate':
            sns.scatterplot(data = df, x = varBias + '_abs_bias', y = varName + '_shap_delta', ax = ax, hue = varHue, s = ms, 
                            palette = ["#65C2A5",'#C7B18A',"#A89BAF",'#F8D347'])
        else:
            sns.scatterplot(data = df, x = varBias + '_abs_bias', y = varName + '_shap_delta', ax = ax, hue = varHue, s = ms)
        r, p = sp.stats.pearsonr(df[varBias+'_abs_bias'].values, df[varName + '_shap_delta'].values) 
    # add lm line
    ax.text(.95, .95, 'r = {:.3f}\np = {:.3f}'.format(r, p) if p >= 0.001 else 'r={:.3f}\np < 0.001'.format(r),
            transform=ax.transAxes, size = fontsize, ha = 'right', va = 'top')
    # legend
    if varHue == 'climate':
        sns.move_legend(ax, loc = 'lower left', markerscale = 1.5, title = None)
    else:
        sns.move_legend(ax, loc = 'lower left', title = varHue, markerscale = 1.5, title_fontsize = fontsize)

    ax.set_xlabel('Bias of %s (mm)'%varBias, fontsize = fontsize)
    ax.set_ylabel(ylabel, fontsize = fontsize)

############################################################################################
import string
# for model in np.unique(df_all.model.values):
    
model = 'ACCESS-ESM1-5'
df = shap_test(model)
fig, axes = plt.subplots(2, 3, figsize = (8, 5), )
# plt.subplots_adjust(hspace = .2, wspace = .2)

plot_pdp(axes[0,0], df, climateVariable, legend = True, ms = 5, fontsize = 8)
plot_pdp(axes[0,1], df, 'ai', legend = False, ms = 5, fontsize = 8)
plot_pdp(axes[0,2], df, 'slope', legend = False, ms = 5, fontsize = 8)

scatter_shap_bias(axes[1,0], df, climateVariable, 5, 'climate', ms = 8, fontsize = 8)
scatter_shap_bias(axes[1,1], df, climateVariable, 'ai', 'ai', ms = 8, fontsize = 8)
scatter_shap_bias(axes[1,2], df, climateVariable, 'slope', 'slope', ms = 8, fontsize = 8)

sns.move_legend(axes[1,1], loc = 'lower left', title = 'Aridity index')
sns.move_legend(axes[1,2], loc = 'lower left', title = 'Slope')
axes[1,1].set_xlim(-250, 110)

fig.tight_layout()

for i,ax0 in enumerate(axes.ravel()):
    ax0.text(-.2, 1.03, string.ascii_letters[i], weight = 'bold', ha = 'right', va = 'bottom', fontsize = 10, transform = ax0.transAxes)

fig.savefig('../picture/fig3.png', dpi = 600)