"""Semi-knockoffs (arXiv:2601.23124v1), implemented from the paper's algorithms.

Algorithm 1 (SKO-Wcx), verbatim:
    Fit  nu_j  ~= E[X^j | X^{-j}]
    Fit  rho_j ~= E[X^j | X^{-j}, y]
    eps_{j,1} = X^j - nu_j(X^{-j})
    eps_{j,2} = X^j - rho_j(X^{-j}, y)
    draw permutations pi_{j,1}, pi_{j,2} of {1..n}
    Xt1_i = nu_j(X^{-j}_i)      + eps_{j,1, pi1(i)}
    Xt2_i = rho_j(X^{-j}_i, y_i) + eps_{j,2, pi2(i)}
    Wilcoxon paired test between {l(m(Xt1_i), y_i)} and {l(m(Xt2_i), y_i)}

Why it is valid (Section 3.1): under H0, rho_j(X^{-j}, y) = E[X^j | X^{-j}, y]
= E[X^j | X^{-j}] = nu_j(X^{-j}), so the two copies are drawn from the *same*
distribution and the paired differences are symmetric about zero — hence a
nonparametric paired test (sign / Wilcoxon) is exact in finite samples. A t-test
is explicitly NOT valid here because the variance vanishes under the null.

Under H1, y carries information about X^j, so rho_j predicts X^j better than
nu_j, copy 2 is perturbed less, and its loss is lower.

FDR (Section 3.2 + Eq. 1): the signed statistic W_j has a sign that is a fair
coin under the null, so the standard knockoff+ threshold applies:
    T_q = min{ t in |W| : (1 + #{W_j <= -t}) / (#{W_j >= t} v 1) <= q }
    S   = { j : W_j >= T_q }
"""
import numpy as np
from scipy.stats import wilcoxon
from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor
from sklearn.linear_model import RidgeCV
from sklearn.neural_network import MLPRegressor


def _fit_predict(model, Z, target):
    model.fit(Z, target)
    return model.predict(Z)


def _regressor(kind, seed):
    if kind == "ridge":
        return RidgeCV(alphas=np.logspace(-3, 3, 13))
    if kind == "rf":
        return RandomForestRegressor(n_estimators=100, random_state=seed, n_jobs=-1)
    if kind == "gb":
        return GradientBoostingRegressor(random_state=seed)
    if kind == "nn":
        return MLPRegressor(hidden_layer_sizes=(64, 32), max_iter=600,
                            random_state=seed)
    raise ValueError(kind)


def semi_knockoff_losses(X, y, j, model, rng, nuisance="ridge", seed=0):
    """Return the two paired loss vectors for feature j (Algorithm 1)."""
    n = len(y)
    Xmj = np.delete(X, j, axis=1)
    xj = X[:, j]

    nu = _fit_predict(_regressor(nuisance, seed), Xmj, xj)
    rho = _fit_predict(_regressor(nuisance, seed),
                       np.column_stack([Xmj, y]), xj)

    e1 = xj - nu
    e2 = xj - rho
    p1 = rng.permutation(n)
    p2 = rng.permutation(n)

    Xt1 = X.copy(); Xt1[:, j] = nu + e1[p1]
    Xt2 = X.copy(); Xt2[:, j] = rho + e2[p2]

    l1 = (model.predict(Xt1) - y) ** 2
    l2 = (model.predict(Xt2) - y) ** 2
    return l1, l2


def sko_pvalue(X, y, j, model, rng, nuisance="ridge", seed=0):
    """Algorithm 1: one-sided Wilcoxon signed-rank p-value for H0: j is null."""
    l1, l2 = semi_knockoff_losses(X, y, j, model, rng, nuisance, seed)
    d = l1 - l2
    if np.allclose(d, 0):
        return 1.0
    # H1: copy 2 (y-aware) has the SMALLER loss, i.e. d > 0
    return float(wilcoxon(d, alternative="greater", zero_method="zsplit").pvalue)


def sko_statistic(X, y, j, model, rng, nuisance="ridge", seed=0):
    """Signed statistic W_j: mean loss of copy 1 minus copy 2.

    Under H0 the two copies are exchangeable, so sign(W_j) is a fair coin.
    Under H1, W_j > 0.
    """
    l1, l2 = semi_knockoff_losses(X, y, j, model, rng, nuisance, seed)
    return float(l1.mean() - l2.mean())


def knockoff_plus_threshold(W, q):
    """Eq. (1): the knockoff+ data-dependent threshold."""
    W = np.asarray(W, float)
    cand = np.sort(np.unique(np.abs(W[W != 0])))
    for t in cand:
        num = 1 + np.sum(W <= -t)
        den = max(np.sum(W >= t), 1)
        if num / den <= q:
            return float(t)
    return float("inf")


def sko_select(X, y, model, rng, q=0.2, nuisance="ridge", seed=0):
    """Algorithm 3: FDR-controlled selection at level q."""
    W = np.array([sko_statistic(X, y, j, model, rng, nuisance, seed)
                  for j in range(X.shape[1])])
    T = knockoff_plus_threshold(W, q)
    return np.where(W >= T)[0], W, T


def fdp_power(selected, nonnull):
    sel = set(int(s) for s in selected)
    nn = set(int(s) for s in nonnull)
    if not sel:
        return 0.0, 0.0
    fdp = len(sel - nn) / len(sel)
    power = len(sel & nn) / max(len(nn), 1)
    return fdp, power


def ar1_design(n, p, rho, rng):
    """AR(1) correlated Gaussian design."""
    Z = rng.standard_normal((n, p))
    X = np.empty_like(Z)
    X[:, 0] = Z[:, 0]
    s = np.sqrt(1 - rho ** 2)
    for j in range(1, p):
        X[:, j] = rho * X[:, j - 1] + s * Z[:, j]
    return X