# Import dependencies
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import time

from shapleyx import rshdmr
from shapleyx.utilities.mc_shapley import GaussianCopulaUniform

from importlib.metadata import version
print(f"Running on ShapleyX v{version('shapleyx')}")

# Set random seed for reproducibility
np.random.seed(123)

def ishigami_vectorised(x1, x2, x3, a=7.0, b=0.1):
    """Ishigami function for vectorised evaluation."""
    return np.sin(x1) + a * np.sin(x2)**2 + b * x3**4 * np.sin(x1)


d = 3
N_train = 256

# Generate training data with a Sobol' sequence
from scipy.stats import qmc
sampler = qmc.Sobol(d, scramble=True, seed=42)
X_raw = sampler.random_base2(int(np.log2(N_train)))
X_train = -np.pi + 2 * np.pi * X_raw  # scale to [-pi, pi]
Y_train = ishigami_vectorised(X_train[:, 0], X_train[:, 1], X_train[:, 2])

dataframe = pd.DataFrame(X_train, columns=['X1', 'X2', 'X3'])
dataframe['Y'] = Y_train
print(f'{len(dataframe)} training samples generated')

model = rshdmr(
    dataframe,
    polys=[10, 5],
    n_iter=300,
    method='ard_cv',
    cv_method='bayesian',
    cv_tol=0.01,
    resampling=True,
    number_of_resamples=500,
)

sob, shap, total = model.run_all()

# Display the HDMR-based Shapley effects for independent inputs
print("\nRS-HDMR Shapley effects (independent inputs):")
shap

print("=" * 60)
print("Permutation method sweep")
print("=" * 60)

for rho in rho_values:
    corr = np.array([
        [1.0, 0.0, rho],
        [0.0, 1.0, 0.0],
        [rho, 0.0, 1.0],
    ])
    t0 = time.time()
    mc_df = model.get_mc_shapley(
        corr=corr,
        N=N_perm,
        method='permutation',
        n_perm=n_perm_val,
        B=B_perm,
        alpha=alpha_96,
        random_state=42,
        progress=False,
    )
    elapsed = time.time() - t0
    mc_df['rho'] = rho
    results_perm.append(mc_df)
    print(f"  rho = {rho:+.1f}  |  "
          f"X1={mc_df['effect'].iloc[0]:.3f}  "
          f"X2={mc_df['effect'].iloc[1]:.3f}  "
          f"X3={mc_df['effect'].iloc[2]:.3f}  "
          f"({elapsed:.0f}s)")

# Concatenate all permutation results
df_perm = pd.concat(results_perm, ignore_index=True)

print("\n" + "=" * 60)
print("Exhaustive method (anchor values)")
print("=" * 60)

for rho in exhaustive_rho:
    corr = np.array([
        [1.0, 0.0, rho],
        [0.0, 1.0, 0.0],
        [rho, 0.0, 1.0],
    ])
    t0 = time.time()
    mc_df = model.get_mc_shapley(
        corr=corr,
        N=N_exh,
        method='exhaustive',
        B=B_exh,
        alpha=alpha_96,
        random_state=42,
        progress=False,
    )
    elapsed = time.time() - t0
    mc_df['rho'] = rho
    results_exh.append(mc_df)
    print(f"  rho = {rho:+.1f}  |  "
          f"X1={mc_df['effect'].iloc[0]:.3f}  "
          f"X2={mc_df['effect'].iloc[1]:.3f}  "
          f"X3={mc_df['effect'].iloc[2]:.3f}  "
          f"({elapsed:.0f}s)")

df_exh = pd.concat(results_exh, ignore_index=True)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6.5))

variables = ['X1', 'X2', 'X3']
colors = ['#2166ac', '#b2182b', '#4daf4a']
markers = ['s', 'o', 'D']

# --- Left panel: permutation sweep with 96 % CI bands ---
for i, var in enumerate(variables):
    mask = df_perm['variable'] == var
    rho_vals = df_perm.loc[mask, 'rho'].values
    eff = df_perm.loc[mask, 'effect'].values
    lo = df_perm.loc[mask, 'lower'].values
    hi = df_perm.loc[mask, 'upper'].values

    # Sort by rho for clean lines
    order = np.argsort(rho_vals)
    ax1.plot(rho_vals[order], eff[order],
             color=colors[i], marker=markers[i], markersize=7,
             linewidth=1.8, label=var)
    ax1.fill_between(rho_vals[order], lo[order], hi[order],
                     color=colors[i], alpha=0.12)

ax1.axhline(0, color='gray', linewidth=0.5, linestyle='--')
ax1.set_xlabel(r'Correlation $\rho$ between $X_1$ and $X_3$')
ax1.set_ylabel('Normalised Shapley Effect')
ax1.set_title('Permutation method (96 % CI bands)')
ax1.legend(loc='best')
ax1.set_ylim(0, None)
ax1.grid(True, alpha=0.3)

# --- Right panel: exhaustive vs permutation at anchor rho ---
bar_width = 0.08
x_pos = np.arange(len(variables))

for j, rho in enumerate(exhaustive_rho):
    # Exhaustive (bars)
    exh_sub = df_exh[df_exh['rho'] == rho]
    offset_exh = (j - 3/3) * bar_width * 1.5
    ax2.bar(
        x_pos + offset_exh, exh_sub['effect'].values, bar_width,
        yerr=[exh_sub['effect'].values - exh_sub['lower'].values,
              exh_sub['upper'].values - exh_sub['effect'].values],
        capsize=3, color=colors[j % len(colors)], alpha=0.75,
        label=f'Exh. ρ={rho:+.1f}',
    )
    # Permutation (scatter markers)
    perm_sub = df_perm[df_perm['rho'] == rho]
    ax2.scatter(
        x_pos + offset_exh, perm_sub['effect'].values,
        marker='x', s=60, color='black', linewidths=1.5, zorder=5,
    )

ax2.set_xticks(x_pos)
ax2.set_xticklabels(variables)
ax2.set_xlabel('Input Variable')
ax2.set_ylabel('Normalised Shapley Effect')
ax2.set_title('Exhaustive (bars) vs Permutation (×) — 96 % CI')
ax2.legend(loc='upper right', fontsize=8)
ax2.set_ylim(0, None)
ax2.grid(True, alpha=0.3, axis='y')

fig.suptitle(
    'Shapley Effects for the Ishigami Function under Correlation\n'
    '(cf. Iooss & Prieur, 2019, Figure 5)',
    fontsize=13, fontweight='bold', y=1.01
)
plt.tight_layout()
plt.show()

rows = []
for rho in exhaustive_rho:
    exh_sub = df_exh[df_exh['rho'] == rho].set_index('variable')
    perm_sub = df_perm[df_perm['rho'] == rho].set_index('variable')
    for var in variables:
        rows.append({
            'rho': rho,
            'Variable': var,
            'Exhaustive': f"{exh_sub.loc[var, 'effect']:.4f}",
            'Permutation': f"{perm_sub.loc[var, 'effect']:.4f}",
            'Diff': f"{abs(exh_sub.loc[var, 'effect'] - perm_sub.loc[var, 'effect']):.4f}",
            'Exh. 96% CI': f"[{exh_sub.loc[var, 'lower']:.3f}, {exh_sub.loc[var, 'upper']:.3f}]",
            'Perm. 96% CI': f"[{perm_sub.loc[var, 'lower']:.3f}, {perm_sub.loc[var, 'upper']:.3f}]",
        })

pd.DataFrame(rows)

print("Exhaustive sweep for Sobol indices (all rho values)")
print("=" * 60)

results_sobol = []

for rho in rho_values:
    corr = np.array([
        [1.0, 0.0, rho],
        [0.0, 1.0, 0.0],
        [rho, 0.0, 1.0],
    ])
    t0 = time.time()
    mc_df = model.get_mc_shapley(
        corr=corr,
        N=5000,
        method='exhaustive',
        B=100,
        alpha=0.04,
        random_state=42,
        progress=False,
    )
    elapsed = time.time() - t0
    mc_df["rho"] = rho
    results_sobol.append(mc_df)
    print(f"  rho = {rho:+.1f}  |  "
          f"S: X1={mc_df["sobol_first"].iloc[0]:.3f}  "
          f"X2={mc_df["sobol_first"].iloc[1]:.3f}  "
          f"X3={mc_df["sobol_first"].iloc[2]:.3f}  "
          f"  T: X1={mc_df["sobol_total"].iloc[0]:.3f}  "
          f"X2={mc_df["sobol_total"].iloc[1]:.3f}  "
          f"X3={mc_df["sobol_total"].iloc[2]:.3f}  "
          f"({elapsed:.0f}s)")

df_sobol = pd.concat(results_sobol, ignore_index=True)

fig, axes = plt.subplots(1, 3, figsize=(18, 5.5))

# Define which CI columns to use for each panel
panel_info = [
    (axes[0], "sobol_first", "First-order Sobol $S_i$",
     "sobol_first_lower", "sobol_first_upper"),
    (axes[1], "effect", "Shapley Effect",
     "lower", "upper"),
    (axes[2], "sobol_total", "Total-order Sobol $T_i$",
     "sobol_total_lower", "sobol_total_upper"),
]

for ax, col, title, ci_lo, ci_hi in panel_info:
    colors = ["#2166ac", "#b2182b", "#4daf4a"]
    markers = ["s", "o", "D"]
    for i, var in enumerate(variables):
        mask = df_sobol["variable"] == var
        rho_vals = df_sobol.loc[mask, "rho"].values
        vals = df_sobol.loc[mask, col].values
        lo = df_sobol.loc[mask, ci_lo].values
        hi = df_sobol.loc[mask, ci_hi].values

        order = np.argsort(rho_vals)
        ax.plot(rho_vals[order], vals[order],
                color=colors[i], marker=markers[i], markersize=7,
                linewidth=1.8, label=var)
        ax.fill_between(rho_vals[order], lo[order], hi[order],
                        color=colors[i], alpha=0.12)

    ax.axhline(0, color="gray", linewidth=0.5, linestyle="--")
    ax.set_xlabel(r"Correlation $\rho$ between $X_1$ and $X_3$")
    ax.set_ylabel("Sensitivity Index")
    ax.set_title(title)
    if ax == axes[0]:
        ax.legend(loc="best")
    ax.set_ylim(0, None)
    ax.grid(True, alpha=0.3)

fig.suptitle(
    "Sobol and Shapley Indices vs Correlation — Ishigami Function\n"
    "(Exhaustive method, 96 % CI bands, cf. Iooss & Prieur 2019)",
    fontsize=12, fontweight="bold", y=1.02
)
plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(8, 5))

for i, var in enumerate(variables):
    mask = df_supp['variable'] == var
    rho_vals = df_supp.loc[mask, 'rho'].values
    eff = df_supp.loc[mask, 'effect'].values
    lo = df_supp.loc[mask, 'lower'].values
    hi = df_supp.loc[mask, 'upper'].values
    order = np.argsort(rho_vals)
    ax.plot(rho_vals[order], eff[order],
            color=colors[i], marker=markers[i], markersize=7,
            linewidth=1.8, label=var)
    ax.fill_between(rho_vals[order], lo[order], hi[order],
                    color=colors[i], alpha=0.12)

ax.axhline(0, color='gray', linewidth=0.5, linestyle='--')
ax.set_xlabel(r'Correlation $\rho$ between $X_1$ and $X_2$')
ax.set_ylabel('Normalised Shapley Effect')
ax.set_title(
    'Shapley Effects — correlation between $X_1$ and $X_2$\n'
    '(Permutation method, 96 % CI)',
    fontsize=11
)
ax.legend(loc='best')
ax.set_ylim(0, None)
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

%load_ext watermark
%watermark -n -u -v -iv -w