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

from shapleyx import rshdmr
from shapleyx.utilities.mc_shapley import (
    GaussianCopulaUniform,
    MCShapley,
    coalitions_up_to_k,
)

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

def owen_product(m, d, tau, mu):
    """Generate samples from the Owen product function."""
    sampler = qmc.Sobol(d, scramble=True, seed=123)
    X = sampler.random_base2(m)
    Y = np.ones(2**m)
    for dim in range(d):
        g = np.sqrt(12) * (X[:, 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([X, Y]), columns=cols + ['Y'])


d = 6
tau = [1.0, 1.0, 0.5, 0.5, 0.25, 0.25]
mu  = [1.0] * d

df = owen_product(m=9, d=d, tau=tau, mu=mu)
print(f"{len(df)} training samples")

# Train RS-HDMR surrogate — polys=[8,6,4] has 4th-order interactions
model = rshdmr(
    df,
    polys=[8, 6, 4, 4],
    n_iter=100,
    method='ard_cv',
    cv_tol=0.005,
)
sob, shap, total = model.run_all()
print(f"\nInteraction order from polys: {len(model.polys)} → k_max={len(model.polys)}")

print("Subset counts for d = 6:")
print(f"  Full enumeration (k_max=None): {len(coalitions_up_to_k(d, None))} subsets")
for kmax in [1, 2, 3, 4, 5]:
    n = len(coalitions_up_to_k(d, kmax))
    pct = 100 * n / len(coalitions_up_to_k(d, None))
    print(f"  k_max={kmax}: {n:3d} subsets ({pct:5.1f}% of full)")

# Auto-detected k_max=3 (from polys=[8,6,4])
t0 = time.time()
mc_auto = model.get_mc_shapley(N=3000, method='permutation', B=100, random_state=42, progress=True)
t_auto = time.time() - t0

# Full enumeration (k_max=None)
t0 = time.time()
mc_full = model.get_mc_shapley(N=3000, method='permutation', B=100,
                              k_max=2, random_state=42, progress=True)
t_full = time.time() - t0

print(f"Auto-detected (k_max=3): {t_auto:.1f}s")
print(f"Full enumeration:         {t_full:.1f}s")
print(f"Speedup: {t_full/t_auto:.1f}x")

# Compare results
comparison = pd.DataFrame({
    'Variable': mc_full['variable'],
    'Full (63 subsets)': mc_full['effect'].round(4),
    'k_max=3 (42 subsets, approximate)': mc_auto['effect'].round(4),
    'Abs. Diff': np.abs(mc_full['effect'] - mc_auto['effect']).round(4),
})
comparison

shap

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

ax.bar(x - w/2, mc_full['effect'], w,
       yerr=[mc_full['effect'] - mc_full['lower'],
             mc_full['upper'] - mc_full['effect']],
       capsize=4, color='steelblue', alpha=0.85,
       label='Full (63 subsets)')
ax.bar(x + w/2, mc_auto['effect'], w,
       yerr=[mc_auto['effect'] - mc_auto['lower'],
             mc_auto['upper'] - mc_auto['effect']],
       capsize=4, color='darkorange', alpha=0.85,
       label='k_max=3 (42 subsets, approximate)')

ax.set_xlabel('Input Variable')
ax.set_ylabel('Shapley Effect')
ax.set_title('Coalition Truncation: Full vs k_max=3\n(Owen product function, d=6, RS-HDMR polys=[8,6,4,4])')
ax.set_xticks(x)
ax.set_xticklabels(mc_full['variable'])
ax.legend(fontsize=9)
ax.set_ylim(0, None)
plt.tight_layout()
plt.show()

# Compare different k_max values
results = {}
for kmax in [1, 2, 3, None]:
    label = f'k_max={kmax}' if kmax else 'Full'
    mc = model.get_mc_shapley(N=2000, method='exhaustive',
                             k_max=kmax, random_state=42)
    results[label] = mc['effect'].values

comparison_df = pd.DataFrame({
    'Variable': mc_full['variable'],
    'k_max=1 (1st order)': results['k_max=1'].round(4),
    'k_max=2 (pairwise)': results['k_max=2'].round(4),
    'k_max=3 (3rd, approx)': results['k_max=3'].round(4),
    'Full (all orders)': results['Full'].round(4),
})
comparison_df

fig, ax = plt.subplots(figsize=(11, 5))
x = np.arange(d)
w = 0.18
labels = ['k_max=1', 'k_max=2', 'k_max=3', 'Full']
colors = ['#92c5de', '#4393c3', '#2166ac', '#053061']

for i, (label, color) in enumerate(zip(labels, colors)):
    ax.bar(x + (i - 1.5) * w, results[label], w,
           color=color, alpha=0.85, label=label)

ax.set_xlabel('Input Variable')
ax.set_ylabel('Shapley Effect')
ax.set_title('Convergence of Shapley Effects with Increasing k_max\n(Owen product function, d=6)')
ax.set_xticks(x)
ax.set_xticklabels(mc_full['variable'])
ax.legend(fontsize=9)
ax.set_ylim(0, None)
plt.tight_layout()
plt.show()

d_vals = np.arange(4, 21)
full_counts = [len(coalitions_up_to_k(d, None)) for d in d_vals]
k1_counts = [len(coalitions_up_to_k(d, 1)) for d in d_vals]
k2_counts = [len(coalitions_up_to_k(d, 2)) for d in d_vals]
k3_counts = [len(coalitions_up_to_k(d, 3)) for d in d_vals]

fig, ax = plt.subplots(figsize=(9, 5))
ax.semilogy(d_vals, full_counts, 'o-', color='black', label='Full (2^d - 1)')
ax.semilogy(d_vals, k3_counts, 's-', color='#2166ac', label='k_max=3')
ax.semilogy(d_vals, k2_counts, 'D-', color='#4393c3', label='k_max=2')
ax.semilogy(d_vals, k1_counts, '^-', color='#92c5de', label='k_max=1')

ax.set_xlabel('Input Dimension d')
ax.set_ylabel('Number of Subsets (log scale)')
ax.set_title('Coalition Truncation: Subset Count vs Dimension')
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

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