#!/usr/bin/env python
"""
Numerical experiments for the leverage correction technical note.

Reproduces all tables and numerical claims in leverage_correction.tex.
Run from the repository root:
    uv run --extra dev python docs/leverage_correction/experiments/run_all.py
"""
import jax
import jax.numpy as jnp
from jax.numpy.linalg import cholesky
from jax.scipy.linalg import solve_triangular
from scipy.stats import spearmanr
import numpy as np
import mellon

jax.config.update("jax_enable_x64", True)


def section(title):
    print()
    print("=" * 64)
    print(title)
    print("=" * 64)


# ------------------------------------------------------------------ #
#  Table 2: Cholesky leverage matches direct inverse (Section 5.1)
# ------------------------------------------------------------------ #
section("Table 2: Cholesky vs direct leverage")

print(f"  {'n':>4} {'sigma':>6} {'max_err':>12} {'tr(H)':>8}")
print(f"  {'-'*4} {'-'*6} {'-'*12} {'-'*8}")

for n in [30, 50, 100]:
    for sigma in [0.5, 1.0, 2.0]:
        key = jax.random.PRNGKey(42)
        k1, k2 = jax.random.split(key)
        X = jax.random.normal(k1, (n, 2))
        y = jnp.sum(jnp.sin(X), axis=1) + sigma * jax.random.normal(k2, (n,))

        est = mellon.FunctionEstimator(sigma=sigma, n_landmarks=0)
        est.fit(X, y)

        h = est.predict.leverage(X, sigma=sigma)

        K = est.predict.cov_func(X, X)
        H_explicit = K @ jnp.linalg.inv(K + sigma**2 * jnp.eye(n))
        h_explicit = jnp.diag(H_explicit)

        max_err = float(jnp.max(jnp.abs(h - h_explicit)))
        trace = float(jnp.sum(h))
        print(f"  {n:4d} {sigma:6.1f} {max_err:12.2e} {trace:8.1f}")


# ------------------------------------------------------------------ #
#  Trace identity: tr(H) = sum lambda_j/(lambda_j + sigma^2)
# ------------------------------------------------------------------ #
section("Trace identity verification")

for n in [50, 100]:
    for sigma in [0.5, 1.0, 2.0]:
        key = jax.random.PRNGKey(42)
        X = jax.random.normal(key, (n, 2))
        y = jnp.zeros(n)

        est = mellon.FunctionEstimator(sigma=sigma, n_landmarks=0)
        est.fit(X, y)

        h = est.predict.leverage(X, sigma=sigma)
        K = est.predict.cov_func(X, X)
        eigvals = jnp.linalg.eigvalsh(K)
        trace_eig = float(jnp.sum(eigvals / (eigvals + sigma**2)))
        trace_h = float(jnp.sum(h))
        print(f"  n={n:3d}, sigma={sigma:.1f}: "
              f"tr(H)_leverage={trace_h:.4f}, "
              f"tr(H)_eigenvalue={trace_eig:.4f}, "
              f"diff={abs(trace_h - trace_eig):.2e}")


# ------------------------------------------------------------------ #
#  LOO residual identity (Section 5.3)
# ------------------------------------------------------------------ #
section("LOO residual identity: r/(1-h) = alpha/diag(Kinv)")

n, sigma = 30, 1.0
key = jax.random.PRNGKey(42)
k1, k2 = jax.random.split(key)
X = jax.random.normal(k1, (n, 2))
y = jnp.sum(jnp.sin(X), axis=1) + sigma * jax.random.normal(k2, (n,))

est = mellon.FunctionEstimator(sigma=sigma, n_landmarks=0)
est.fit(X, y)
pred = est.predict(X)
r = y - pred
h = est.predict.leverage(X, sigma=sigma)
r_corrected = r / (1 - h)

K = est.predict.cov_func(X, X)
Kinv = jnp.linalg.inv(K + sigma**2 * jnp.eye(n))
alpha = Kinv @ (y - est.predict.mu)
r_loo_analytic = alpha / jnp.diag(Kinv)

max_err = float(jnp.max(jnp.abs(r_corrected - r_loo_analytic)))
print(f"  n={n}, sigma={sigma}: max|r/(1-h) - alpha/diag(Kinv)| = {max_err:.2e}")


# ------------------------------------------------------------------ #
#  Table 3: Sparse vs full leverage (Section 5.4)
# ------------------------------------------------------------------ #
section("Table 3: Sparse vs full leverage correlation")

print(f"  {'n':>4} {'m':>4} {'sigma':>6} {'Spearman':>10} {'max_diff':>10}")
print(f"  {'-'*4} {'-'*4} {'-'*6} {'-'*10} {'-'*10}")

for n, m in [(50, 20), (50, 30), (100, 30), (100, 50)]:
    for sigma in [0.5, 1.0, 2.0]:
        key = jax.random.PRNGKey(42)
        k1, k2 = jax.random.split(key)
        X = jax.random.normal(k1, (n, 2))
        y = jnp.sum(jnp.sin(X), axis=1) + sigma * jax.random.normal(k2, (n,))

        est_full = mellon.FunctionEstimator(sigma=sigma, n_landmarks=0)
        est_full.fit(X, y)
        h_full = est_full.predict.leverage(X, sigma=sigma)

        est_sparse = mellon.FunctionEstimator(sigma=sigma, n_landmarks=m)
        est_sparse.fit(X, y)
        h_sparse = est_sparse.predict.leverage(X, sigma=sigma)

        corr, _ = spearmanr(h_full, h_sparse)
        max_diff = float(jnp.max(jnp.abs(h_full - h_sparse)))
        print(f"  {n:4d} {m:4d} {sigma:6.1f} {corr:10.4f} {max_diff:10.4f}")


# ------------------------------------------------------------------ #
#  Monte Carlo bias decomposition (Section 5.5)
# ------------------------------------------------------------------ #
section("Monte Carlo bias decomposition (n=50, sigma=1)")

n, sigma = 50, 1.0
key = jax.random.PRNGKey(42)
X = jax.random.normal(key, (n, 2))

est = mellon.FunctionEstimator(sigma=sigma, n_landmarks=0)
est.fit(X, jnp.zeros(n))
K = est.predict.cov_func(X, X)
H = K @ jnp.linalg.inv(K + sigma**2 * jnp.eye(n))
h = jnp.diag(H)
I_minus_H = jnp.eye(n) - H

signal = jnp.sin(X[:, 0])
n_mc = 10_000
mean_r2 = jnp.zeros(n)

for i in range(n_mc):
    eps = sigma * jax.random.normal(jax.random.PRNGKey(i), (n,))
    r_i = I_minus_H @ (signal + eps)
    mean_r2 = mean_r2 + r_i**2
mean_r2 = mean_r2 / n_mc

noise_part = sigma**2 * jnp.diag(I_minus_H @ I_minus_H.T)
signal_misfit = (I_minus_H @ signal)**2
exact_total = noise_part + signal_misfit
hc3_approx = sigma**2 * (1 - h)**2

print(f"  mean|MC - sigma^2(1-h)^2|  = {float(jnp.mean(jnp.abs(mean_r2 - hc3_approx))):.4f}")
print(f"  mean|MC - exact decomp|    = {float(jnp.mean(jnp.abs(mean_r2 - exact_total))):.4f}")
print(f"  signal misfit fraction     = {float(jnp.mean(signal_misfit) / jnp.mean(exact_total)):.2%}")
print(f"  noise fraction             = {float(jnp.mean(noise_part) / jnp.mean(exact_total)):.2%}")


# ------------------------------------------------------------------ #
#  Table 4: Homoscedastic noise recovery (Section 5.6)
# ------------------------------------------------------------------ #
section("Table 4: Homoscedastic noise recovery (n=200, sigma_GP=1.0)")

n = 200
key = jax.random.PRNGKey(42)
k1, k2 = jax.random.split(key)
X = jax.random.normal(k1, (n, 2))

print(f"  {'true_s2':>8} {'mean_obs_var':>12} {'std':>8} {'ratio':>8}")
print(f"  {'-'*8} {'-'*12} {'-'*8} {'-'*8}")

for true_sigma in [0.5, 1.0, 2.0, 5.0]:
    y = jnp.sin(X[:, 0]) + true_sigma * jax.random.normal(k2, (n,))
    est = mellon.FunctionEstimator(sigma=1.0, n_landmarks=0, obs_variance=True)
    est.fit(X, y)
    obs_var = est.predict.obs_variance(X)
    mean_v = float(jnp.mean(obs_var))
    std_v = float(jnp.std(obs_var))
    print(f"  {true_sigma**2:8.2f} {mean_v:12.3f} {std_v:8.3f} {mean_v / true_sigma**2:8.3f}")


# ------------------------------------------------------------------ #
#  Smoothing reduces CV (Section 5.7)
# ------------------------------------------------------------------ #
section("Smoothing reduces CV (n=200, true_sigma=2)")

n = 200
true_sigma = 2.0
key = jax.random.PRNGKey(42)
k1, k2 = jax.random.split(key)
X = jax.random.normal(k1, (n, 2))
y = jnp.sin(X[:, 0]) + true_sigma * jax.random.normal(k2, (n,))

est = mellon.FunctionEstimator(sigma=true_sigma, n_landmarks=0, obs_variance=True)
est.fit(X, y)
h = est.predict.leverage(X, sigma=true_sigma)
r = y - est.predict(X)
r2_corrected = r**2 / (1 - h)**2
obs_var_smoothed = est.predict.obs_variance(X)

cv_raw = float(jnp.std(r2_corrected) / jnp.mean(r2_corrected))
cv_smooth = float(jnp.std(obs_var_smoothed) / jnp.mean(obs_var_smoothed))

print(f"  Raw corrected r^2:  mean={float(jnp.mean(r2_corrected)):.3f}, CV={cv_raw:.3f}")
print(f"  GP-smoothed:        mean={float(jnp.mean(obs_var_smoothed)):.3f}, CV={cv_smooth:.3f}")
print(f"  Theoretical chi2(1) CV: {np.sqrt(2):.3f}")
print(f"  CV reduction: {(1 - cv_smooth/cv_raw)*100:.0f}%")


# ------------------------------------------------------------------ #
#  Table 1: Regime analysis (Section 4)
# ------------------------------------------------------------------ #
section("Table 1: Regime analysis (n=50, d=3)")

n = 50
key = jax.random.PRNGKey(42)
k1, k2 = jax.random.split(key)
X = jax.random.normal(k1, (n, 3))
y = jnp.sin(X[:, 0]) + 0.5 * jnp.cos(X[:, 1]) + jax.random.normal(k2, (n,))

print(f"  {'ls_factor':>10} {'sigma':>6} {'h_max':>8} {'max_corr':>10} {'tr(H)':>8}")
print(f"  {'-'*10} {'-'*6} {'-'*8} {'-'*10} {'-'*8}")

for ls_factor in [0.5, 1.0, 5.0, 10.0]:
    for sigma in [0.1, 0.5, 1.0, 2.0]:
        est = mellon.FunctionEstimator(
            sigma=sigma, n_landmarks=0, ls_factor=ls_factor
        )
        est.fit(X, y)
        h = est.predict.leverage(X, sigma=sigma)
        h_max = float(h.max())
        correction_max = 1 / (1 - h_max)**2
        trace = float(jnp.sum(h))
        print(f"  {ls_factor:10.1f} {sigma:6.1f} {h_max:8.4f} "
              f"{correction_max:10.2f} {trace:8.1f}")


# ------------------------------------------------------------------ #
#  Numerical stability (Section 3.1 footnote)
# ------------------------------------------------------------------ #
section("Cholesky numerical stability across sigma range")

print(f"  {'n':>4} {'sigma':>8} {'max_err':>12}")
print(f"  {'-'*4} {'-'*8} {'-'*12}")

for n in [50, 200]:
    for sigma in [0.1, 1.0, 10.0]:
        key = jax.random.PRNGKey(42)
        X = jax.random.normal(key, (n, 2))
        est = mellon.FunctionEstimator(sigma=sigma, n_landmarks=0)
        est.fit(X, jnp.zeros(n))
        K = est.predict.cov_func(X, X)

        h_direct = jnp.diag(K @ jnp.linalg.inv(K + sigma**2 * jnp.eye(n)))

        L = cholesky(K + sigma**2 * jnp.eye(n) + 1e-6 * jnp.eye(n))
        Linv = solve_triangular(L, jnp.eye(n), lower=True)
        h_chol = 1 - sigma**2 * jnp.sum(Linv**2, axis=0)

        max_err = float(jnp.max(jnp.abs(h_direct - h_chol)))
        print(f"  {n:4d} {sigma:8.1f} {max_err:12.2e}")


# ------------------------------------------------------------------ #
#  Full vs sparse obs_variance consistency (Section 5.8)
# ------------------------------------------------------------------ #
section("Full vs sparse obs_variance consistency")

print(f"  {'n':>4} {'m':>4} {'sigma':>6} {'lev_rho':>10} {'var_rho':>10}")
print(f"  {'-'*4} {'-'*4} {'-'*6} {'-'*10} {'-'*10}")

for n, m in [(50, 20), (100, 30), (100, 50)]:
    for sigma in [0.5, 1.0, 2.0]:
        key = jax.random.PRNGKey(42)
        k1, k2 = jax.random.split(key)
        X = jax.random.normal(k1, (n, 2))
        y = jnp.sum(jnp.sin(X), axis=1) + sigma * jax.random.normal(k2, (n,))

        est_full = mellon.FunctionEstimator(sigma=sigma, n_landmarks=0, obs_variance=True)
        est_full.fit(X, y)
        h_full = est_full.predict.leverage(X, sigma=sigma)
        var_full = est_full.predict.obs_variance(X)

        est_sparse = mellon.FunctionEstimator(sigma=sigma, n_landmarks=m, obs_variance=True)
        est_sparse.fit(X, y)
        h_sparse = est_sparse.predict.leverage(X, sigma=sigma)
        var_sparse = est_sparse.predict.obs_variance(X)

        lev_rho, _ = spearmanr(h_full, h_sparse)
        var_rho, _ = spearmanr(var_full, var_sparse)
        print(f"  {n:4d} {m:4d} {sigma:6.1f} {lev_rho:10.4f} {var_rho:10.4f}")


print()
print("All experiments completed successfully.")
