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

from shapleyx.utilities.mc_shapley import (
    MultivariateNormal,
    GaussianCopulaUniform,
    shapley_effects,
)

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

d = 10

# Variable labels and uniform ranges
labels = ['Sw', 'Wfw', 'A', 'Lambda', 'q', 'l', 'tc', 'Nz', 'Wdg', 'Wp']
descriptions = [
    'Wing area (ft²)',
    'Fuel weight in wing (lb)',
    'Aspect ratio',
    'Quarter-chord sweep (deg)',
    'Dynamic pressure at cruise (lb/ft²)',
    'Taper ratio',
    'Airfoil thickness/chord ratio',
    'Ultimate load factor',
    'Flight design gross weight (lb)',
    'Paint weight (lb/ft²)',
]

# Uniform bounds for each variable
lows = np.array([150, 220, 6, -10, 16, 0.5, 0.08, 2.5, 1700, 0.025])
highs = np.array([200, 300, 10, 10, 45, 1.0, 0.18, 6.0, 2500, 0.08])

print(f"{'Variable':>8s}  {'Range':>25s}  {'Description':s}")
print("-" * 75)
for i, (name, desc) in enumerate(zip(labels, descriptions)):
    print(f"{name:>8s}  [{lows[i]:6g}, {highs[i]:6g}]   {desc}")

# --- Independent inputs (baseline) ---
corr_independent = np.eye(d)
joint_ind = GaussianCopulaUniform(lows, highs, corr_independent)

# --- Correlated inputs (extension) ---
corr_correlated = np.eye(d)
corr_correlated[0, 2] = corr_correlated[2, 0] = 0.6   # Sw ↔ A
corr_correlated[0, 8] = corr_correlated[8, 0] = 0.7   # Sw ↔ Wdg
corr_correlated[2, 5] = corr_correlated[5, 2] = -0.4  # A ↔ λ
corr_correlated[3, 4] = corr_correlated[4, 3] = 0.3   # Λ ↔ q

joint_corr = GaussianCopulaUniform(lows, highs, corr_correlated)

print("GaussianCopulaUniform distributions ready.")
print(f"Dimension: {d}")
print(f"Baseline: independent inputs")
print(f"Extension: 4 correlation pairs among geometric/aero params")

def wing_weight(x):
    """Wing weight (lb) from Forrester et al. (2008).

    x = [Sw, Wfw, A, Lambda, q, l, tc, Nz, Wdg, Wp]
    """
    Sw, Wfw, A, Lambda, q, l, tc, Nz, Wdg, Wp = x

    # Convert sweep angle to radians for cos()
    cos_L = np.cos(np.deg2rad(Lambda))

    term1 = 0.036 * Sw**0.758 * Wfw**0.0035
    term2 = (A / cos_L**2)**0.6
    term3 = q**0.006 * l**0.04
    term4 = (100 * tc / cos_L)**(-0.3)
    term5 = (Nz * Wdg)**0.49
    term6 = Sw * Wp

    return term1 * term2 * term3 * term4 * term5 + term6


def wing_weight_batch(X):
    """Vectorised wing weight for batch evaluation."""
    Sw  = X[:, 0]
    Wfw = X[:, 1]
    A   = X[:, 2]
    Lambda = X[:, 3]
    q   = X[:, 4]
    l   = X[:, 5]
    tc  = X[:, 6]
    Nz  = X[:, 7]
    Wdg = X[:, 8]
    Wp  = X[:, 9]

    cos_L = np.cos(np.deg2rad(Lambda))

    term1 = 0.036 * Sw**0.758 * Wfw**0.0035
    term2 = (A / cos_L**2)**0.6
    term3 = q**0.006 * l**0.04
    term4 = (100 * tc / cos_L)**(-0.3)
    term5 = (Nz * Wdg)**0.49
    term6 = Sw * Wp

    return term1 * term2 * term3 * term4 * term5 + term6


# Quick sanity check at nominal values
x_nom = [175, 260, 8, 0, 30, 0.75, 0.13, 4.25, 2100, 0.0525]
W_nom = wing_weight(x_nom)
print(f"Nominal wing weight: {W_nom:.1f} lb")

# Quick histogram of wing weight values
X_check = joint_ind.sample_joint(50000)
W_sample = wing_weight_batch(X_check)

fig, ax = plt.subplots(figsize=(8, 4))
ax.hist(W_sample, bins=80, density=True, color='steelblue', alpha=0.85)
ax.set_xlabel('Wing Weight (lb)')
ax.set_ylabel('Density')
ax.set_title('Wing Weight Distribution (Independent Inputs)')
plt.tight_layout()
plt.show()

print(f"Wing weight range: [{W_sample.min():.0f}, {W_sample.max():.0f}] lb")
print(f"Mean: {W_sample.mean():.0f} lb, Std: {W_sample.std():.0f} lb")
print(f"CoV: {W_sample.std() / W_sample.mean():.3f}")

eff, sh, var, lower, upper = shapley_effects(
    wing_weight,
    joint_ind,
    N=10000,
    method='exhaustive',
    B=200,
    alpha=0.05,
    predict_batch=wing_weight_batch,
    random_state=42,
    progress=True,
)

results = pd.DataFrame({
    'Variable': labels,
    'Effect': eff,
    'Lower': lower,
    'Upper': upper,
})
results

# Generate training data — Sobol' samples over full ranges
from scipy.stats.qmc import Sobol

N_train = 2048
sampler = Sobol(d=d, scramble=True, seed=123)
X_sobol = sampler.random(N_train)
# Scale to input ranges
X_train = lows + (highs - lows) * X_sobol
Y_train = wing_weight_batch(X_train)

# Build DataFrame for RS-HDMR
df_train = pd.DataFrame(X_train, columns=labels)
df_train['Y'] = Y_train

print(f"{len(df_train)} training samples")
print(f"Y range: [{Y_train.min():.1f}, {Y_train.max():.1f}] lb")
print(f"Mean: {Y_train.mean():.1f}, Std: {Y_train.std():.1f} lb")

from shapleyx import rshdmr

# Fit RS-HDMR surrogate — up to 3rd-order interactions
model = rshdmr(
    df_train,
    polys=[8, 6, 4],
    n_iter=200,
    method='ard_cv',
    cv_method='bayesian',
    cv_tol=0.005,
)

sob_surr, shap_surr, total_surr = model.run_all()

# MC Shapley using the surrogate with independent distribution
mc_surrogate = model.get_mc_shapley(
    joint=joint_ind,
    N=10000,
    method='exhaustive',
    B=200,
    alpha=0.05,
    random_state=42,
    progress=True,
)
mc_surrogate

# Published first-order Sobol indices (OpenTURNS, Saltelli 2002)
# From: https://openturns.github.io/openturns/latest
ref_sobol_first = np.array([
    0.130315,    # Sw
    2.94e-06,    # Wfw
    0.228153,    # A
    0.0,         # Lambda
    8.25e-05,    # q
    0.001803,    # l (taper ratio)
    0.135002,    # tc
    0.412794,    # Nz
    0.088332,    # Wdg
    0.003516,    # Wp
])

# Our RS-HDMR Sobol first-order indices
our_sobol_mc = mc_surrogate['sobol_first'].values

validation = pd.DataFrame({
    'Variable': labels,
    'MC Sobol S_i': our_sobol_mc,
    'RS-HDMR Sobol S_i': filtered_df['index'],
    'Published S_i (OpenTURNS)': ref_sobol_first,
})
validation

# MC Shapley for correlated inputs (via RS-HDMR surrogate)
mc_corr = model.get_mc_shapley(
    joint=joint_corr,
    N=10000,
    method='exhaustive',
    B=200,
    alpha=0.05,
    random_state=42,
    progress=True,
)
mc_corr

comparison = pd.DataFrame({
    'Variable': labels,
    'Shapley (analytical, indep)': eff.round(4),
    'Shapley (surrogate, indep)': mc_surrogate['effect'].values.round(4),
    'Shapley (surrogate, corr)': mc_corr['effect'].values.round(4),
    'Sobol S_i (pub. OT ref)': ref_sobol_first.round(4),
})
comparison

# Bar chart comparison
fig, ax = plt.subplots(figsize=(16, 7))
bar_width = 0.18
x = np.arange(d)

ax.bar(x - 1.5*bar_width, eff, bar_width,
       yerr=[eff - lower, upper - eff],
       capsize=2, color='steelblue', alpha=0.85,
       label='Shapley (analytical, independent)')
ax.bar(x - 0.5*bar_width, mc_surrogate['effect'].values, bar_width,
       yerr=[mc_surrogate['effect'].values - mc_surrogate['lower'].values,
             mc_surrogate['upper'].values - mc_surrogate['effect'].values],
       capsize=2, color='darkorange', alpha=0.85,
       label='Shapley (RS-HDMR, independent)')
ax.bar(x + 0.5*bar_width, ref_sobol_first, bar_width,
       color='crimson', alpha=0.75,
       label='Sobol $S_i$ (OpenTURNS reference)')

ax.bar(x + 1.5*bar_width, mc_corr['effect'].values, bar_width,
       yerr=[mc_corr['effect'].values - mc_corr['lower'].values,
             mc_corr['upper'].values - mc_corr['effect'].values],
       capsize=2, color='seagreen', alpha=0.85,
       label='Shapley (RS-HDMR, correlated)')

ax.set_xlabel('Input Variable')
ax.set_ylabel('Sensitivity Index')
ax.set_title('Wing Weight Function — Shapley Effects vs. Sobol Indices\\n'
             'Independent vs. Correlated Inputs')
ax.set_xticks(x)
ax.set_xticklabels(labels)
ax.legend(fontsize=9, ncol=2)
ax.set_ylim(0, None)
plt.tight_layout()
plt.show()

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