"""Graph-aware uncertainty: component bootstrap, jackknife, Benjamini-Hochberg.

Algorithms copied verbatim from ppi_benchmark_controls.py v1.0.0 (lines 718-773).
Do not edit the algorithms without re-running the golden tests."""
from __future__ import annotations

import numpy as np
import pandas as pd
from scipy import stats
from sklearn.metrics import roc_auc_score

from .audit import connected_components


def component_groups(test_pairs: pd.DataFrame) -> list[np.ndarray]:
    """Row indices grouped by connected component of the test-pair graph."""
    comps = connected_components(test_pairs)
    node_to_comp = {n: i for i, comp in enumerate(comps) for n in comp}
    comp_of_row = test_pairs["protein_a"].map(node_to_comp).to_numpy()
    return [np.where(comp_of_row == c)[0] for c in range(len(comps))]


def component_bootstrap_ci(y, p, groups, n_boot: int, seed: int, null: float = 0.5) -> dict:
    """Percentile bootstrap resampling whole components with replacement."""
    rng = np.random.default_rng(seed)
    vals = []
    for _ in range(n_boot):
        idx = np.concatenate([groups[k] for k in rng.integers(0, len(groups), len(groups))])
        if len(np.unique(y[idx])) == 2:
            vals.append(roc_auc_score(y[idx], p[idx]))
    vals = np.array(vals)
    if len(vals) == 0:
        return {"boot_low": np.nan, "boot_high": np.nan, "boot_p": np.nan}
    return {"boot_low": np.percentile(vals, 2.5), "boot_high": np.percentile(vals, 97.5),
            "boot_p": min(1.0, 2 * min(np.mean(vals <= null), np.mean(vals >= null)))}


def component_jackknife_ci(y, p, groups, null: float = 0.5) -> dict:
    """Delete-one-component jackknife with a t(G_eff - 1) reference distribution."""
    point = roc_auc_score(y, p)
    all_idx = np.arange(len(y))
    loo = []
    for rows in groups:
        keep = np.setdiff1d(all_idx, rows, assume_unique=True)
        if len(np.unique(y[keep])) == 2:
            loo.append(roc_auc_score(y[keep], p[keep]))
    loo = np.array(loo)
    g = len(loo)
    if g < 2:
        return {"jack_low": np.nan, "jack_high": np.nan, "jack_p": np.nan, "components_used": g}
    se = np.sqrt((g - 1) / g * np.sum((loo - loo.mean()) ** 2))
    t = stats.t.ppf(0.975, df=g - 1)
    pval = (0.0 if point != null else 1.0) if se == 0 else 2 * stats.t.sf(abs((point - null) / se), df=g - 1)
    return {"jack_low": point - t * se, "jack_high": point + t * se, "jack_p": pval, "components_used": g}


def benjamini_hochberg(pvals: list[float]) -> list[float]:
    p = np.asarray(pvals, float)
    out = np.full(len(p), np.nan)
    ok = ~np.isnan(p)
    m = ok.sum()
    if m:
        idx = np.where(ok)[0][np.argsort(p[ok])]
        adj = np.minimum.accumulate((p[idx] * m / np.arange(1, m + 1))[::-1])[::-1]
        out[idx] = np.clip(adj, 0, 1)
    return out.tolist()


# =============================================================================
# 8. Figures