# Import dependencies
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from itertools import combinations
from scipy.stats import qmc

from shapleyx import rshdmr

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

def owen_product(m, d, tau, mu):
    """Generate samples from the Owen product function.

    Uses a scrambled Sobol' sequence (base-2) for quasi-Monte Carlo
    sampling and returns a DataFrame with columns X1...Xd and Y.
    """
    num_samples = 2 ** m
    sampler = qmc.Sobol(d, scramble=True, seed=123)
    samples = sampler.random_base2(m)

    Y = np.ones(num_samples)
    for dim in range(d):
        g = np.sqrt(12) * (samples[:, dim] - 0.5)
        Y *= (mu[dim] + tau[dim] * g)

    cols = [f'X{i+1}' for i in range(d)]
    return pd.DataFrame(
        np.column_stack([samples, Y]),
        columns=cols + ['Y']
    )


def analytical_sigma2(u, tau, mu):
    """Analytical Sobol' index for subset u."""
    d = len(tau)
    prod_tau = np.prod([tau[j] ** 2 for j in u])
    prod_mu = np.prod([mu[j] ** 2 for j in range(d) if j not in u])
    return prod_tau * prod_mu


def analytical_shapley(tau, mu):
    """Compute analytical Shapley effects for the product function.

    For the Owen product function the Shapley value for variable i is
    the sum of all Sobol' indices involving i, each divided by its
    interaction order (Owen 2014).  Normalised by the true total
    variance of the product function.
    """
    d = len(tau)
    all_vars = list(range(d))
    shapleys = np.zeros(d)

    # True total variance: prod(mu_j^2 + tau_j^2) - prod(mu_j^2)
    total_var = (np.prod([mu[j]**2 + tau[j]**2 for j in range(d)])
                 - np.prod([mu[j]**2 for j in range(d)]))

    for r in range(1, d + 1):
        for subset in combinations(all_vars, r):
            sig2 = analytical_sigma2(subset, tau, mu)
            share = sig2 / len(subset)
            for i in subset:
                shapleys[i] += share

    return shapleys / total_var


print("Helper functions defined.")

m = 8
d = 6
tau = [1.0, 1.0, 0.5, 0.5, 0.25, 0.25]
mu  = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]

dataframe = owen_product(m, d, tau, mu)
print(f"{len(dataframe)} samples generated")

model = rshdmr(
    dataframe,
    polys=[8, 6, 4, 4],
    n_iter=100,
    method='ard_cv',
    cv_tol=0.005,
)

sob, shap, total = model.run_all()

shap

# Add analytical column to the Shapley DataFrame
shap['analytical'] = analytical_shapley(tau, mu)
shap

comparison = pd.DataFrame({
    'Variable': shap['label'],
    'RS-HDMR': shap['scaled effect'].round(4),
    'Analytical': shap['analytical'].round(4),
})
comparison['Abs. Difference'] = np.abs(
    comparison['RS-HDMR'] - comparison['Analytical']
).round(4)
comparison

fig, ax = plt.subplots(figsize=(10, 6))
bar_width = 0.3
x = np.arange(d)

ax.bar(x - bar_width/2, shap['scaled effect'], bar_width,
       yerr=[shap['scaled effect'] - shap['lower'],
             shap['upper'] - shap['scaled effect']],
       capsize=5, color='steelblue', alpha=0.85,
       label='RS-HDMR surrogate')
ax.bar(x + bar_width/2, shap['analytical'], bar_width,
       color='darkorange', alpha=0.85,
       label='Analytical (exact)')

ax.set_xlabel('Input Variable')
ax.set_ylabel('Shapley Effect')
ax.set_title('RS-HDMR vs Analytical Shapley Effects\n(Owen product function, d = 6)')
ax.set_xticks(x)
ax.set_xticklabels(shap['label'])
ax.legend()
ax.set_ylim(0, None)
plt.tight_layout()
plt.show()

# First-order Sobol' indices
first_order = sob[~sob['derived_labels'].str.contains('_')].copy()
first_order = first_order.rename(columns={'derived_labels': 'variable'})
first_order[['variable', 'index', 'lower', 'upper']].round(4)

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