"""Benjamini-Hochberg FDR across ablated sections (docs/05-statistics.md).

K sections means K simultaneous comparisons. At K=12 with α=0.05 the chance of
at least one spurious "this section matters" is about 46%, so leaving it
unadjusted would manufacture a finding in roughly half of all runs.

**FDR, not Bonferroni.** The ablation stage is screening for follow-up, so the
error worth controlling is the false-*discovery* rate — what fraction of the
sections I promote to a confirmatory run will be wild geese. Bonferroni
controls the family-wise error rate instead, which is the wrong question here
and at K=15 would leave nothing detectable at any affordable N.

This is one of three layers, and the weakest of them (docs/05-statistics.md):
the output is framed as ranking rather than testing, adjusted here, and then
the survivors get a dedicated confirmatory A/B where a single pre-specified
comparison needs no adjustment at all.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

#: Screening for follow-up tolerates more false discoveries than a
#: confirmatory analysis would, because the confirmatory stage is what
#: filters them. q=0.10 is the figure in docs/05-statistics.md.
DEFAULT_FDR_Q = 0.10


@dataclass(frozen=True)
class FDRResult:
    """Adjusted p-values and which sections survive at `q`."""

    adjusted: tuple[float, ...]
    rejected: tuple[bool, ...]
    q: float

    @property
    def n_discoveries(self) -> int:
        return sum(self.rejected)


def benjamini_hochberg(p_values, q: float = DEFAULT_FDR_Q) -> FDRResult:
    """Adjust `p_values` for multiplicity, controlling FDR at `q`.

    Returns adjusted values in the caller's original order. NaN inputs — a
    section whose interval was withheld, so it has no p to adjust — stay NaN
    and are never counted as discoveries.
    """
    raw = np.asarray(list(p_values), dtype=float)
    adjusted = np.full(raw.shape, np.nan)
    rejected = np.zeros(raw.shape, dtype=bool)

    usable = np.flatnonzero(np.isfinite(raw))
    m = usable.size
    if m == 0:
        return FDRResult(tuple(adjusted), tuple(rejected), q)

    order = usable[np.argsort(raw[usable], kind="stable")]
    ranks = np.arange(1, m + 1)

    # Step up from the largest p: each adjusted value is the running minimum of
    # (m/rank) * p from the top down, which keeps the sequence monotone. Without
    # that, a section could show a smaller adjusted p than one ranked above it.
    scaled = raw[order] * m / ranks
    adjusted[order] = np.minimum.accumulate(scaled[::-1])[::-1].clip(max=1.0)
    rejected[order] = adjusted[order] <= q

    return FDRResult(tuple(adjusted), tuple(rejected), q)
