"""Fresh CPU audit of the Semi-knockoffs Section 5.1 comparison.

This is an independent execution, not a read of the authors' CSV outputs.  It
uses the paper's adjacent-support design (Sigma_ij=0.6**abs(i-j), first
quarter of beta nonzero, gradient boosting) and a separately implemented HRT
holdout.  A second exact construction exercises the masked-correlation
five-permutation assertion.
"""
import json
import sys
from pathlib import Path

import numpy as np
from sklearn.base import clone
from sklearn.ensemble import GradientBoostingRegressor
from scipy.stats import wilcoxon

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from skocore import ar1_design


ALPHA = 0.05
OUT = {"alpha": ALPHA, "design": "independent_cpu_claim5_scope_audit"}


def source_design(n, p, seed):
    rng = np.random.default_rng(seed)
    # ar1_design is the exact Gaussian construction for Sigma_ij=rho^|i-j|.
    X = ar1_design(n, p, 0.6, rng)
    beta = np.zeros(p)
    beta[: p // 4] = rng.uniform(1.0, 2.0, size=p // 4)
    y = np.einsum("ij,j->i", X, beta) + rng.normal(size=n)
    return X, y, rng


def ridge_predict(Z_train, target, Z_eval, alpha=1.0):
    """Closed-form finite Ridge prediction using einsum, avoiding BLAS matmul."""
    gram = np.einsum("ij,ik->jk", Z_train, Z_train)
    rhs = np.einsum("ij,i->j", Z_train, target)
    coef = np.linalg.solve(gram + alpha * np.eye(Z_train.shape[1]), rhs)
    return np.einsum("ij,j->i", Z_eval, coef)


def semi_pvalue(X, y, j, model, rng, seed, n_perm=1):
    """Standalone Algorithm-1 p-value with closed-form Ridge nuisances."""
    Xmj = np.delete(X, j, axis=1)
    xj = X[:, j]
    nu = ridge_predict(Xmj, xj, Xmj)
    rho = ridge_predict(np.column_stack([Xmj, y]), xj,
                        np.column_stack([Xmj, y]))
    e1 = xj - nu
    e2 = xj - rho
    diffs = []
    for _ in range(n_perm):
        x1 = X.copy()
        x2 = X.copy()
        x1[:, j] = nu + e1[rng.permutation(len(y))]
        x2[:, j] = rho + e2[rng.permutation(len(y))]
        diffs.append((model.predict(x1) - y) ** 2 -
                     (model.predict(x2) - y) ** 2)
    d = np.concatenate(diffs)
    if np.allclose(d, 0.0):
        return 1.0
    return float(wilcoxon(d, alternative="greater", zero_method="zsplit").pvalue)


def hrt_pvalue(X, y, j, seed, permutations=100):
    """One HRT p-value with one train/test split and a conditional ridge draw."""
    rng = np.random.default_rng(seed)
    order = rng.permutation(len(y))
    n_train = len(y) // 2
    train, test = order[:n_train], order[n_train:]
    model = GradientBoostingRegressor(random_state=seed).fit(X[train], y[train])
    xm = np.delete(X, j, axis=1)
    mu = ridge_predict(xm[train], X[train, j], xm[test])
    residual = X[train, j] - ridge_predict(xm[train], X[train, j], xm[train])
    observed = np.mean((model.predict(X[test]) - y[test]) ** 2)
    count = 0
    for _ in range(permutations):
        xp = X[test].copy()
        xp[:, j] = mu + rng.choice(residual, size=len(test), replace=True)
        count += float(np.mean((model.predict(xp) - y[test]) ** 2) <= observed)
    return (1.0 + count) / (1.0 + permutations)


def adjacent_grid(ns=(200, 300, 400), reps=15):
    rows = []
    for n in ns:
        sko_alt = []
        sko_null = []
        hrt_alt = []
        hrt_null = []
        for r in range(reps):
            seed = 18000 + 100 * n + r
            X, y, rng = source_design(n, 50, seed)
            model = GradientBoostingRegressor(random_state=seed).fit(X, y)
            # j=0 is in the first-quarter adjacent support; j=30 is null.
            sko_alt.append(semi_pvalue(X, y, 0, model, rng, seed=seed) <= ALPHA)
            sko_null.append(semi_pvalue(X, y, 30, model, rng, seed=seed + 1) <= ALPHA)
            hrt_alt.append(hrt_pvalue(X, y, 0, seed + 2) <= ALPHA)
            hrt_null.append(hrt_pvalue(X, y, 30, seed + 3) <= ALPHA)
        row = {
            "n": n,
            "reps": reps,
            "sko_power": float(np.mean(sko_alt)),
            "sko_type_I": float(np.mean(sko_null)),
            "hrt_power": float(np.mean(hrt_alt)),
            "hrt_type_I": float(np.mean(hrt_null)),
        }
        row["power_gap"] = row["sko_power"] - row["hrt_power"]
        rows.append(row)
        print("adjacent", row, flush=True)
    return rows


def masked_grid(reps=15, amplitudes=(0.25, 0.5, 1.0)):
    rows = []
    for amplitude in amplitudes:
        cells = []
        for r in range(reps):
            seed = 24000 + 100 * int(amplitude * 100) + r
            rng = np.random.default_rng(seed)
            X = ar1_design(300, 50, 0.6, rng)
            signal = 25
            null = signal - 1
            y = amplitude * X[:, signal] + 0.5 * rng.normal(size=300)
            # Replace the adjacent coordinate by a correlated but label-null copy.
            X[:, null] = X[:, signal] + 0.5 * rng.normal(size=300)
            model = GradientBoostingRegressor(random_state=seed).fit(X, y)
            one = semi_pvalue(X, y, signal, model, rng, seed=seed, n_perm=1) <= ALPHA
            five = semi_pvalue(X, y, signal, model, rng, seed=seed + 10,
                               n_perm=5) <= ALPHA
            null_p = semi_pvalue(X, y, null, model, rng, seed=seed + 30,
                                 n_perm=5)
            cells.append({
                "one_reject": bool(one),
                "five_reject": bool(five),
                "five_null_reject": bool(null_p <= ALPHA),
            })
        row = {
            "signal_amplitude": amplitude,
            "reps": reps,
            "single_power": float(np.mean([r["one_reject"] for r in cells])),
            "five_power": float(np.mean([r["five_reject"] for r in cells])),
            "five_null_type_I": float(np.mean([r["five_null_reject"] for r in cells])),
        }
        row["power_gain"] = row["five_power"] - row["single_power"]
        rows.append(row)
        print("masked", row, flush=True)
    out = {"reps": reps, "amplitudes": list(amplitudes), "cells": rows}
    print("masked", out, flush=True)
    return out


def main():
    OUT["adjacent"] = adjacent_grid()
    OUT["masked"] = masked_grid()
    out = Path("outputs/claim5_scope_audit.json")
    out.write_text(json.dumps(OUT, indent=2) + "\n")
    print("saved", out, flush=True)


if __name__ == "__main__":
    main()