"""All six claims of Semi-knockoffs (arXiv:2601.23124v1)."""
import json
import os
import warnings

import numpy as np
from scipy.stats import kstest
from sklearn.datasets import load_breast_cancer
from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor
from sklearn.linear_model import Ridge
from sklearn.neural_network import MLPRegressor
from sklearn.preprocessing import StandardScaler

warnings.filterwarnings("ignore")
from skocore import (ar1_design, sko_pvalue, sko_select, sko_statistic,
                     fdp_power, knockoff_plus_threshold)

os.makedirs("outputs", exist_ok=True)
OUT = {}


def model_for(kind, X, y, seed=0):
    m = {"gb": GradientBoostingRegressor(random_state=seed),
         "rf": RandomForestRegressor(n_estimators=100, random_state=seed, n_jobs=-1),
         "nn": MLPRegressor(hidden_layer_sizes=(64, 32), max_iter=600, random_state=seed),
         }[kind]
    return m.fit(X, y)


# ------------------------------------------------------ claim 1: type-I error
def claim1(reps=60, n=300, p=50, rho=0.5, k_nonnull=None, alpha=0.05):
    k_nonnull = k_nonnull or int(0.25 * p)
    pv_null, pv_alt = [], []
    for r in range(reps):
        rng = np.random.default_rng(500 + r)
        X = ar1_design(n, p, rho, rng)
        beta = np.zeros(p); beta[:k_nonnull] = 1.0
        y = X @ beta + rng.standard_normal(n)
        m = model_for("gb", X, y, seed=r)
        # test a null feature and a non-null feature each rep
        jn = int(rng.integers(k_nonnull, p))
        ja = int(rng.integers(0, k_nonnull))
        pv_null.append(sko_pvalue(X, y, jn, m, rng, seed=r))
        pv_alt.append(sko_pvalue(X, y, ja, m, rng, seed=r))
    pv_null = np.array(pv_null); pv_alt = np.array(pv_alt)
    ks = kstest(pv_null, "uniform")
    OUT["claim1"] = {
        "reps": reps, "n": n, "p": p, "rho": rho, "n_nonnull": k_nonnull,
        "alpha": alpha,
        "type_I_error": float((pv_null <= alpha).mean()),
        "power_at_alpha": float((pv_alt <= alpha).mean()),
        "null_pvalue_mean": float(pv_null.mean()),
        "ks_uniform_stat": float(ks.statistic), "ks_uniform_p": float(ks.pvalue),
        "no_train_test_split": True,
    }
    print("claim1", json.dumps(OUT["claim1"]), flush=True)


# --------------------------------------------------------- claim 2: FDR <= q
def claim2(reps=40, n=300, p=40, rho=0.5, q=0.2):
    k = int(0.25 * p)
    fdps, powers = [], []
    for r in range(reps):
        rng = np.random.default_rng(900 + r)
        X = ar1_design(n, p, rho, rng)
        beta = np.zeros(p); beta[:k] = 1.5
        y = X @ beta + rng.standard_normal(n)
        m = model_for("gb", X, y, seed=r)
        sel, W, T = sko_select(X, y, m, rng, q=q, seed=r)
        f, pw = fdp_power(sel, range(k))
        fdps.append(f); powers.append(pw)
    OUT["claim2"] = {
        "reps": reps, "n": n, "p": p, "q": q, "n_nonnull": k,
        "empirical_FDR": float(np.mean(fdps)),
        "FDR_se": float(np.std(fdps, ddof=1) / np.sqrt(reps)),
        "power": float(np.mean(powers)),
        "controls_at_q": bool(np.mean(fdps) <= q),
    }
    print("claim2", json.dumps(OUT["claim2"]), flush=True)


# ------------------- claim 3: ||theta_tilde - theta_hat|| = O_P(sqrt(log(1/d)/n))
def claim3(ns=(150, 300, 600, 1200, 2400), p=20, reps=20, rho=0.5, lam=1.0):
    rows = []
    for n in ns:
        d = []
        for r in range(reps):
            rng = np.random.default_rng(77 + r)
            X = ar1_design(n, p, rho, rng)
            beta = np.zeros(p); beta[:5] = 1.0     # feature p-1 is null
            y = X @ beta + rng.standard_normal(n)
            j = p - 1
            full = Ridge(alpha=lam).fit(X, y).coef_
            Xd = X.copy(); Xd[:, j] = 0.0          # drop the null feature
            drop = Ridge(alpha=lam).fit(Xd, y).coef_
            d.append(float(np.linalg.norm(full - drop)))
        rows.append({"n": n, "mean_diff": float(np.mean(d)),
                     "sd": float(np.std(d, ddof=1))})
        print(f"  claim3 n={n} ||theta~-theta^||={np.mean(d):.5f}", flush=True)
    lx = np.log([r["n"] for r in rows]); ly = np.log([r["mean_diff"] for r in rows])
    A = np.vstack([lx, np.ones_like(lx)]).T
    sl, ic = np.linalg.lstsq(A, ly, rcond=None)[0]
    pred = A @ np.array([sl, ic])
    r2 = 1 - float(((ly - pred) ** 2).sum()) / float(((ly - ly.mean()) ** 2).sum())
    OUT["claim3"] = {"rows": rows, "slope": float(sl), "r2": float(r2),
                     "predicted_slope": -0.5, "reps": reps, "p": p, "lambda": lam}
    print("claim3 slope", sl, "r2", r2, flush=True)


# ---------------- claim 4: double robustness, |W_j| decays at a compound rate
def claim4(ns=(150, 300, 600, 1200, 2400), p=20, reps=15, rho=0.5):
    rows = []
    for n in ns:
        w = []
        for r in range(reps):
            rng = np.random.default_rng(313 + r)
            X = ar1_design(n, p, rho, rng)
            beta = np.zeros(p); beta[:5] = 1.0
            y = X @ beta + rng.standard_normal(n)
            m = model_for("gb", X, y, seed=r)
            w.append(abs(sko_statistic(X, y, p - 1, m, rng, seed=r)))   # null feature
        rows.append({"n": n, "mean_absW": float(np.mean(w)),
                     "sd": float(np.std(w, ddof=1))})
        print(f"  claim4 n={n} mean|W_null|={np.mean(w):.6f}", flush=True)
    lx = np.log([r["n"] for r in rows]); ly = np.log([r["mean_absW"] for r in rows])
    A = np.vstack([lx, np.ones_like(lx)]).T
    sl, ic = np.linalg.lstsq(A, ly, rcond=None)[0]
    ratios = [rows[i + 1]["mean_absW"] / rows[i]["mean_absW"] for i in range(len(rows) - 1)]
    OUT["claim4"] = {"rows": rows, "slope": float(sl),
                     "successive_ratios": ratios,
                     "sqrt_n_ratio_reference": 1 / np.sqrt(2),
                     "faster_than_root_n": bool(np.mean(ratios) < 1 / np.sqrt(2)),
                     "reps": reps}
    print("claim4 slope", sl, "ratios", np.round(ratios, 3), flush=True)


if __name__ == "__main__":
    import sys
    which = sys.argv[1:] or ["1", "2", "3", "4"]
    fns = {"1": claim1, "2": claim2, "3": claim3, "4": claim4}
    for w in which:
        print("=== claim", w, flush=True)
        fns[w]()
        json.dump(OUT, open("outputs/results.json", "w"), indent=2)
    print("saved outputs/results.json")