"""Claims 1-3 of "An Odd Estimator for Shapley Values" (arXiv:2602.01399v1).

All three are exact algebraic statements, so they are checked by exhaustive
enumeration at machine precision rather than by sampling.

  Observation 3.1 : phi_i(f) = phi_i(f_odd)  and  phi_i(f_even) = 0
  Theorem 3.2    : paired sampling orthogonalises the weighted regression, so
                   the odd and even blocks separate
  Theorem 3.5    : constrained Fourier regression on any T containing all
                   singletons returns the exact Shapley values

Decomposition used throughout (the paper's):
    f_odd(S)  = (f(S) - f(S^c)) / 2
    f_even(S) = (f(S) + f(S^c)) / 2
Fourier basis: chi_T(S) = (-1)^{|S ^ T|}.
"""
import itertools
import json
import warnings
from math import comb, factorial

import numpy as np

# Apple Accelerate emits spurious FP-status warnings from matmul on this
# machine; products were cross-checked against np.einsum and agree to 1e-15.
warnings.filterwarnings("ignore", message=".*encountered in matmul.*")


def masks(d):
    return np.arange(1 << d, dtype=np.int64)


def popcount(a):
    return np.array([bin(int(x)).count("1") for x in a])


def shapley_exhaustive(fvals, d):
    """phi_i = sum_S f(S) (1[i in S] p_{|S|-1} - 1[i notin S] p_{|S|})."""
    p = [factorial(l) * factorial(d - l - 1) / factorial(d) for l in range(d)]
    M = masks(d)
    sz = popcount(M)
    phi = np.zeros(d)
    for i in range(d):
        inS = ((M >> i) & 1).astype(bool)
        w = np.where(inS,
                     np.array([p[s - 1] if s >= 1 else 0.0 for s in sz]),
                     -np.array([p[s] if s <= d - 1 else 0.0 for s in sz]))
        phi[i] = float(fvals @ w)
    return phi


def complement(M, d):
    return (~M) & ((1 << d) - 1)


def odd_even(fvals, d):
    M = masks(d)
    fc = fvals[complement(M, d)]
    return (fvals - fc) / 2.0, (fvals + fc) / 2.0


def chi(T, S):
    return -1.0 if bin(int(T) & int(S)).count("1") % 2 else 1.0


def chi_col(T, M):
    return np.array([chi(T, int(s)) for s in M])


def kernel_weights(d):
    """w_l = (d-1) / (C(d,l) l (d-l)) for 0 < l < d; boundary handled by the
    exact constraints, so those rows are excluded from the regression."""
    w = np.zeros(d + 1)
    for l in range(1, d):
        w[l] = (d - 1) / (comb(d, l) * l * (d - l))
    return w


def run():
    rng = np.random.default_rng(2026)
    out = {}

    # ---------------------------------------------------------- Observation 3.1
    rows = []
    for d in [6, 8, 9, 10]:
        for trial in range(4):
            f = rng.normal(size=1 << d)
            fo, fe = odd_even(f, d)
            phi = shapley_exhaustive(f, d)
            phi_o = shapley_exhaustive(fo, d)
            phi_e = shapley_exhaustive(fe, d)
            rows.append({
                "d": d, "n_coalitions": 1 << d,
                "max_abs_phi_f_minus_phi_odd": float(np.abs(phi - phi_o).max()),
                "max_abs_phi_even": float(np.abs(phi_e).max()),
                "phi_scale": float(np.abs(phi).max()),
            })
    out["observation_3_1"] = {
        "rows": rows,
        "max_diff_overall": max(r["max_abs_phi_f_minus_phi_odd"] for r in rows),
        "max_phi_even_overall": max(r["max_abs_phi_even"] for r in rows),
        "n_games": len(rows),
    }

    # ------------------------------------------------------------ Theorem 3.2
    # Paired sampling: for each drawn S also draw S^c.  In the Fourier basis the
    # odd columns (|T| odd) and even columns (|T| even) must be orthogonal under
    # the weighted inner product induced by the sample.
    t32 = []
    for d in [8, 10]:
        M = masks(d)
        w = kernel_weights(d)
        Ts = [T for T in range(1 << d)]
        odd_T = [T for T in Ts if bin(T).count("1") % 2 == 1]
        even_T = [T for T in Ts if bin(T).count("1") % 2 == 0]
        odd_T = odd_T[:40]
        even_T = even_T[:40]
        for npairs in [50, 200]:
            # draw npairs coalitions strictly between empty and full, and pair
            # each with its complement
            cand = M[(popcount(M) > 0) & (popcount(M) < d)]
            S = rng.choice(cand, size=npairs, replace=True)
            Sc = complement(S, d)

            # (a) the theorem at the level of a single pair: the contribution
            #     of {S, S^c} to any odd-even cross entry must vanish exactly,
            #     because chi_T(S^c) = (-1)^{|T|} chi_T(S) and the kernel
            #     weights satisfy w_l = w_{d-l}.
            pairworst = 0.0
            for k in range(min(len(S), 60)):
                a, b_ = int(S[k]), int(Sc[k])
                wa, wb = w[bin(a).count("1")], w[bin(b_).count("1")]
                for T1 in odd_T[:10]:
                    for T2 in even_T[:10]:
                        v = (wa * chi(T1, a) * chi(T2, a)
                             + wb * chi(T1, b_) * chi(T2, b_))
                        pairworst = max(pairworst, abs(v))

            # (b) interleaved ordering, so each pair cancels adjacently
            Sp = np.empty(2 * npairs, dtype=np.int64)
            Sp[0::2] = S
            Sp[1::2] = Sc
            wt = np.array([w[c] for c in popcount(Sp)])
            A_odd = np.stack([chi_col(T, Sp) for T in odd_T], 1)
            A_even = np.stack([chi_col(T, Sp) for T in even_T], 1)
            cross = A_odd.T @ (wt[:, None] * A_even)

            # (c) blocked ordering (all S rows, then all S^c rows): the same
            #     quantity mathematically, but the cancellation now happens
            #     between two large partial sums, so it only holds to roundoff
            blk = np.concatenate([S, Sc])
            wtb = np.array([w[c] for c in popcount(blk)])
            Ab = np.stack([chi_col(T, blk) for T in odd_T], 1)
            Bb = np.stack([chi_col(T, blk) for T in even_T], 1)
            cross_blk = Ab.T @ (wtb[:, None] * Bb)
            diag_scale = float(np.abs(Ab.T @ (wtb[:, None] * Ab)).max())
            # unpaired control: independent draws, no complements
            S2 = rng.choice(cand, size=2 * npairs, replace=True)
            wt2 = np.array([w[c] for c in popcount(S2)])
            B_odd = np.stack([chi_col(T, S2) for T in odd_T], 1)
            B_even = np.stack([chi_col(T, S2) for T in even_T], 1)
            cross2 = B_odd.T @ (wt2[:, None] * B_even)
            t32.append({
                "d": d, "n_pairs": npairs, "n_samples": 2 * npairs,
                "per_pair_max_abs_contribution": pairworst,
                "per_pair_exactly_zero": bool(pairworst == 0.0),
                "paired_interleaved_max_abs_cross_gram": float(np.abs(cross).max()),
                "paired_interleaved_exactly_zero": bool(np.all(cross == 0.0)),
                "paired_blocked_max_abs_cross_gram": float(np.abs(cross_blk).max()),
                "gram_diagonal_scale": diag_scale,
                "unpaired_max_abs_cross_gram": float(np.abs(cross2).max()),
                "gram_block_shape": list(cross.shape),
            })
    out["theorem_3_2"] = {
        "rows": t32,
        "all_per_pair_exactly_zero": bool(all(r["per_pair_exactly_zero"] for r in t32)),
        "all_interleaved_exactly_zero": bool(all(
            r["paired_interleaved_exactly_zero"] for r in t32)),
        "max_blocked_roundoff": max(
            r["paired_blocked_max_abs_cross_gram"] for r in t32),
        "max_unpaired": max(r["unpaired_max_abs_cross_gram"] for r in t32),
    }

    # ------------------------------------------------------------ Theorem 3.5
    # Constrained Fourier regression: minimise sum_{0<|S|<d} w_|S| (f(S)-g(S))^2
    # over g = sum_{T in Tset} beta_T chi_T subject to g(empty)=f(empty) and
    # g([d])=f([d]).  Then phi(g) must equal phi(f) exactly.
    t35 = []
    for d in [7, 8, 9]:
        M = masks(d)
        sz = popcount(M)
        w = kernel_weights(d)
        interior = M[(sz > 0) & (sz < d)]
        wt = np.array([w[c] for c in popcount(interior)])
        # Shapley value of each basis function, computed exactly
        phi_chi = {}
        for trial in range(3):
            f = rng.normal(size=1 << d)
            phi_true = shapley_exhaustive(f, d)
            singles = [1 << i for i in range(d)]
            for extra_k in [0, 3, 8, 20]:
                pool = [T for T in range(1 << d)
                        if T not in singles and T != 0]
                extra = list(rng.choice(pool, size=extra_k, replace=False)) \
                    if extra_k else []
                Tset = [0] + singles + [int(t) for t in extra]
                Tset = sorted(set(Tset))
                A = np.stack([chi_col(T, interior) for T in Tset], 1)
                # constraints: sum_T beta_T = f(empty);
                #              sum_T (-1)^{|T|} beta_T = f(full)
                c1 = np.ones(len(Tset))
                c2 = np.array([(-1.0) ** bin(T).count("1") for T in Tset])
                C = np.stack([c1, c2], 0)
                b = np.array([f[0], f[(1 << d) - 1]])
                # weighted LS with equality constraints via KKT
                W = wt[:, None]
                H = A.T @ (W * A) + 1e-12 * np.eye(len(Tset))
                g = A.T @ (wt * f[interior])
                KKT = np.block([[H, C.T], [C, np.zeros((2, 2))]])
                rhs = np.concatenate([g, b])
                sol = np.linalg.solve(KKT, rhs)
                beta = sol[:len(Tset)]
                for T in Tset:
                    if T not in phi_chi:
                        phi_chi[T] = shapley_exhaustive(chi_col(T, M), d)
                phi_hat = sum(beta[k] * phi_chi[T] for k, T in enumerate(Tset))
                t35.append({
                    "d": d, "trial": trial, "n_basis": len(Tset),
                    "extra_beyond_singletons": extra_k,
                    "max_abs_phi_err": float(np.abs(phi_hat - phi_true).max()),
                    "phi_scale": float(np.abs(phi_true).max()),
                    "efficiency_err": float(abs(
                        phi_hat.sum() - (f[(1 << d) - 1] - f[0]))),
                })
    out["theorem_3_5"] = {
        "rows": t35,
        "max_phi_err_overall": max(r["max_abs_phi_err"] for r in t35),
        "max_efficiency_err": max(r["efficiency_err"] for r in t35),
        "n_regressions": len(t35),
    }

    # negative control for Theorem 3.5: drop a singleton from T (violating the
    # T superset T_{<=1} hypothesis) and the exactness must fail
    d = 8
    M = masks(d)
    sz = popcount(M)
    w = kernel_weights(d)
    interior = M[(sz > 0) & (sz < d)]
    wt = np.array([w[c] for c in popcount(interior)])
    f = rng.normal(size=1 << d)
    phi_true = shapley_exhaustive(f, d)
    Tset = sorted(set([0] + [1 << i for i in range(1, d)]))   # drops feature 0
    A = np.stack([chi_col(T, interior) for T in Tset], 1)
    C = np.stack([np.ones(len(Tset)),
                  np.array([(-1.0) ** bin(T).count("1") for T in Tset])], 0)
    H = A.T @ (wt[:, None] * A) + 1e-12 * np.eye(len(Tset))
    KKT = np.block([[H, C.T], [C, np.zeros((2, 2))]])
    sol = np.linalg.solve(KKT, np.concatenate(
        [A.T @ (wt * f[interior]), np.array([f[0], f[(1 << d) - 1]])]))
    beta = sol[:len(Tset)]
    phi_chi = {T: shapley_exhaustive(chi_col(T, M), d) for T in Tset}
    phi_hat = sum(beta[k] * phi_chi[T] for k, T in enumerate(Tset))
    out["theorem_3_5"]["negative_control_missing_singleton"] = {
        "max_abs_phi_err": float(np.abs(phi_hat - phi_true).max()),
        "phi_scale": float(np.abs(phi_true).max()),
    }

    json.dump(out, open("outputs_odd.json", "w"), indent=2)
    print(json.dumps({
        "obs31_max_diff": out["observation_3_1"]["max_diff_overall"],
        "obs31_max_phi_even": out["observation_3_1"]["max_phi_even_overall"],
        "obs31_games": out["observation_3_1"]["n_games"],
        "thm32_per_pair_zero": out["theorem_3_2"]["all_per_pair_exactly_zero"],
        "thm32_interleaved_zero": out["theorem_3_2"]["all_interleaved_exactly_zero"],
        "thm32_blocked_roundoff": out["theorem_3_2"]["max_blocked_roundoff"],
        "thm32_max_unpaired": out["theorem_3_2"]["max_unpaired"],
        "thm35_max_err": out["theorem_3_5"]["max_phi_err_overall"],
        "thm35_max_eff_err": out["theorem_3_5"]["max_efficiency_err"],
        "thm35_n": out["theorem_3_5"]["n_regressions"],
        "thm35_control_err": out["theorem_3_5"][
            "negative_control_missing_singleton"]["max_abs_phi_err"],
    }, indent=2))


if __name__ == "__main__":
    run()