"""Per-cell latency verdicts: A/A floor, bootstrap CI, exact permutation p, BH-FDR.

A "cell" is one (scenario, operation) compared between two subjects in one

lane invocation. The old devloop gate applied ONE threshold (2x an A/A floor

measured on the micro profile) to every cell, so a high-variance cell could

publish an "improvement" of a binary over itself. Here every cell is judged on

its own evidence:

1. Per-cell A/A floor: an A/A control subject (the baseline binary under a

   second subject name, interleaved in the same randomized runs) gives each

   cell its own null. The floor is the largest |relative median delta| inside

   the A/A bootstrap interval, clamped below by ``MIN_FLOOR_PCT``.

2. Effect CI: percentile bootstrap of the relative median delta

   (target / baseline - 1). Runs share indices across subjects (interleaved),

   so resampling is paired by run when both arms have the same run set.

3. p-value: exact two-sided permutation test of the Mann-Whitney rank

   statistic (all C(n1+n2, n1) relabelings; seeded Monte Carlo above

   ``EXACT_LIMIT``).

4. Benjamini-Hochberg over every tested cell in the comparison family.

A cell is a CHANGE only when BH rejects AND the CI excludes zero AND the

effect exceeds the cell's A/A floor (itself the outer A/A CI bound) AND the absolute delta clears the minimum

milliseconds. A cell whose CI and floor say "change" but which BH does not

confirm is UNCONFIRMED: never published as a win or a regression, reported

with the repetitions needed to confirm. Missing A/A evidence or fewer than

``MIN_SAMPLES`` per arm is UNMEASURED, never a pass.

"""

from __future__ import annotations

import itertools

import math

import random

import statistics

from typing import Any, Callable, Iterable

from oakbench.stats import benjamini_hochberg_detail

CELL_STATS_VERSION = "2026-09-27.2"

MIN_FLOOR_PCT = 3.0

MIN_SAMPLES = 5

SCREEN_P = 0.2

EXACT_LIMIT = 20_000

DEFAULT_FDR = 0.05

DEFAULT_CONFIDENCE = 0.95

DEFAULT_BOOTSTRAP = 2000

CellKey = tuple[str, str]

def samples_by_cell(

    rows: Iterable[dict[str, Any]],

    subject: str,

    value_of: Callable[[dict[str, Any]], Any],

    *,

    include: Callable[[dict[str, Any]], bool] | None = None,

) -> dict[CellKey, dict[int, float]]:

    """{(scenario, operation): {run: value}} for successful measured rows."""

    cells: dict[CellKey, dict[int, float]] = {}

    for row in rows:

        if row.get("subject") != subject or row.get("skipped"):

            continue

        if row.get("returncode") not in (0, None) or (include is not None and not include(row)):

            continue

        value = value_of(row)

        if type(value) not in (int, float) or not math.isfinite(value):

            continue

        key = (str(row.get("scenario")), str(row.get("operation")))

        run = row.get("run")

        run_key = run if type(run) is int else len(cells.get(key, {}))

        cells.setdefault(key, {})[run_key] = float(value)

    return cells

def relative_delta_pct(baseline: list[float], target: list[float]) -> float | None:

    base = statistics.median(baseline)

    if base <= 0:

        return None

    return (statistics.median(target) / base - 1.0) * 100.0

def bootstrap_relative_ci(

    baseline: dict[int, float],

    target: dict[int, float],

    *,

    samples: int = DEFAULT_BOOTSTRAP,

    confidence: float = DEFAULT_CONFIDENCE,

    seed: int = 20260927,

) -> tuple[float, float] | None:

    rng = random.Random(seed)

    paired = sorted(baseline) == sorted(target)

    runs = sorted(baseline)

    base_values = list(baseline.values())

    target_values = list(target.values())

    estimates = []

    for _ in range(samples):

        if paired:

            chosen = [runs[rng.randrange(len(runs))] for _ in runs]

            b = [baseline[run] for run in chosen]

            t = [target[run] for run in chosen]

        else:

            b = [base_values[rng.randrange(len(base_values))] for _ in base_values]

            t = [target_values[rng.randrange(len(target_values))] for _ in target_values]

        value = relative_delta_pct(b, t)

        if value is not None:

            estimates.append(value)

    if not estimates:

        return None

    estimates.sort()

    alpha = 1.0 - confidence

    low = estimates[max(0, math.floor(alpha / 2 * len(estimates)))]

    high = estimates[min(len(estimates) - 1, math.ceil((1 - alpha / 2) * len(estimates)) - 1)]

    return low, high

def _mann_whitney_u(baseline: list[float], target: list[float]) -> float:

    """U statistic of target over baseline (ties count one half)."""

    return sum(1.0 if t > b else 0.5 if t == b else 0.0 for t in target for b in baseline)

def permutation_p_value(baseline: list[float], target: list[float], *, seed: int = 20260927,

                        exact_limit: int = EXACT_LIMIT) -> float:

    """Exact two-sided permutation p for the Mann-Whitney rank statistic.

    Rank-based, so one stalled sample cannot manufacture significance; exact

    over all C(n1+n2, n1) relabelings (min p = 2/252 at 5 vs 5), seeded Monte

    Carlo above ``exact_limit``.

    """

    pooled = list(baseline) + list(target)

    n1 = len(baseline)

    centre = n1 * len(target) / 2.0

    observed = abs(_mann_whitney_u(baseline, target) - centre)

    tolerance = 1e-9

    total = math.comb(len(pooled), n1)

    extreme = 0

    if total <= exact_limit:

        indices = range(len(pooled))

        for chosen in itertools.combinations(indices, n1):

            picked = set(chosen)

            a = [pooled[i] for i in chosen]

            b = [pooled[i] for i in indices if i not in picked]

            if abs(_mann_whitney_u(a, b) - centre) >= observed - tolerance:

                extreme += 1

        return extreme / total

    rng = random.Random(seed)

    draws = exact_limit

    for _ in range(draws):

        shuffled = pooled[:]

        rng.shuffle(shuffled)

        if abs(_mann_whitney_u(shuffled[:n1], shuffled[n1:]) - centre) >= observed - tolerance:

            extreme += 1

    return (extreme + 1) / (draws + 1)

def min_exact_p(n1: int, n2: int) -> float:

    """Smallest two-sided exact permutation p-value reachable at n1 vs n2."""

    return min(1.0, 2.0 / math.comb(n1 + n2, n1))

def runs_for_bh(family_size: int, fdr: float, start: int = MIN_SAMPLES) -> int:

    """Runs per arm at which ONE cell can be confirmed by BH in a family of m.

    A lone change is BH-significant only if its p <= fdr / m, and the exact

    permutation p can never be below 2 / C(2n, n).

    """

    n = max(2, start)

    while min_exact_p(n, n) > fdr / max(1, family_size) and n < 200:

        n += 1

    return n

def required_runs(item: dict[str, Any], family_size: int, fdr: float) -> int | None:

    """Runs per arm to confirm this cell's observed effect: the larger of the

    BH feasibility bound and the n that shrinks the CI half-width inside the

    effect's margin over the floor (half-width scales as 1/sqrt(n))."""

    delta, floor = item.get("delta_pct"), item.get("floor_pct")

    low, high = item.get("ci_low_pct"), item.get("ci_high_pct")

    n = min(item.get("baseline_n") or 0, item.get("target_n") or 0)

    if None in (delta, floor, low, high) or n < 2:

        return None

    if abs(delta) <= floor:

        return None  # the observed effect is inside the noise floor; no n confirms it

    half = (high - low) / 2.0

    n_ci = math.ceil(n * (half / abs(delta)) ** 2) if half > 0 else n

    return max(n + 1, n_ci, runs_for_bh(family_size, fdr, n))

def minimum_detectable_effect_pct(item: dict[str, Any]) -> float | None:

    """Smallest true |effect| this cell could have confirmed at its n: the

    point effect must exceed the floor AND the CI (half-width h) must exclude

    zero, so MDE ~= max(floor, h)."""

    floor, low, high = item.get("floor_pct"), item.get("ci_low_pct"), item.get("ci_high_pct")

    if None in (floor, low, high):

        return None

    return max(floor, (high - low) / 2.0)

def screen_candidates(cells: dict[CellKey, dict[str, Any]], *, target_mde_pct: float,

                      screen_p: float = SCREEN_P) -> list[CellKey]:

    """Cells worth an independent confirmation run: anything suggestive

    (p <= screen_p, CI excluding zero, a stage-1 change or unconfirmed

    candidate) or too imprecise to rule out ``target_mde_pct``."""

    chosen = []

    for key, item in cells.items():

        if item.get("p_value") is None:

            continue

        excludes_zero = item["ci_low_pct"] > 0 or item["ci_high_pct"] < 0

        mde = item.get("mde_pct")

        if (item["verdict"] != "no_change" or item["p_value"] <= screen_p or excludes_zero

                or (mde is not None and mde > target_mde_pct)):

            chosen.append(key)

    return chosen

def aa_floor_pct(baseline: dict[int, float], control: dict[int, float], **kwargs: Any) -> float | None:

    """Per-cell A/A floor: largest |delta| the identical binary shows here."""

    if len(baseline) < MIN_SAMPLES or len(control) < MIN_SAMPLES:

        return None

    point = relative_delta_pct(list(baseline.values()), list(control.values()))

    interval = bootstrap_relative_ci(baseline, control, **kwargs)

    if point is None or interval is None:

        return None

    return max(MIN_FLOOR_PCT, abs(point), abs(interval[0]), abs(interval[1]))

def cell_floors(rows: list[dict[str, Any]], baseline: str, control: str,

                value_of: Callable[[dict[str, Any]], Any], **kwargs: Any) -> dict[CellKey, float]:

    base = samples_by_cell(rows, baseline, value_of)

    other = samples_by_cell(rows, control, value_of)

    floors = {}

    for key in sorted(set(base) & set(other)):

        floor = aa_floor_pct(base[key], other[key], **kwargs)

        if floor is not None:

            floors[key] = floor

    return floors

def classify_cells(

    rows: list[dict[str, Any]],

    target: str,

    baseline: str,

    floors: dict[CellKey, float],

    *,

    value_of: Callable[[dict[str, Any]], Any] = lambda row: row.get("elapsed_ms"),

    cells: Iterable[CellKey] | None = None,

    min_delta_ms: float = 0.0,

    fdr: float = DEFAULT_FDR,

    confidence: float = DEFAULT_CONFIDENCE,

    bootstrap_samples: int = DEFAULT_BOOTSTRAP,

    seed: int = 20260927,

) -> dict[CellKey, dict[str, Any]]:

    """Per-cell verdict dicts for one comparison family (see module doc)."""

    base_cells = samples_by_cell(rows, baseline, value_of)

    target_cells = samples_by_cell(rows, target, value_of)

    keys = sorted(set(cells) if cells is not None else (set(base_cells) | set(target_cells)))

    results: dict[CellKey, dict[str, Any]] = {}

    tested: list[CellKey] = []

    p_values: list[float] = []

    for key in keys:

        base = base_cells.get(key, {})

        cand = target_cells.get(key, {})

        item: dict[str, Any] = {

            "cell_stats_version": CELL_STATS_VERSION, "scenario": key[0], "operation": key[1],

            "baseline_n": len(base), "target_n": len(cand), "floor_pct": floors.get(key),

            "baseline_median": statistics.median(base.values()) if base else None,

            "target_median": statistics.median(cand.values()) if cand else None,

            "delta_pct": None, "ci_low_pct": None, "ci_high_pct": None, "p_value": None,

            "fdr_rejected": None, "verdict": "unmeasured", "reason": None,

        }

        results[key] = item

        if len(base) < MIN_SAMPLES or len(cand) < MIN_SAMPLES:

            item["reason"] = f"fewer than {MIN_SAMPLES} measured runs per arm"

            continue

        delta = relative_delta_pct(list(base.values()), list(cand.values()))

        interval = bootstrap_relative_ci(base, cand, samples=bootstrap_samples, confidence=confidence, seed=seed)

        if delta is None or interval is None:

            item["reason"] = "baseline median is zero"

            continue

        item.update(delta_pct=delta, ci_low_pct=interval[0], ci_high_pct=interval[1],

                    p_value=permutation_p_value(list(base.values()), list(cand.values()), seed=seed))

        tested.append(key)

        p_values.append(item["p_value"])

    bh = benjamini_hochberg_detail(p_values, fdr=fdr)

    for key, rejected in zip(tested, bh["rejected"]):

        item = results[key]

        item["fdr_rejected"] = rejected

        floor = item["floor_pct"]

        excludes_zero = item["ci_low_pct"] > 0 or item["ci_high_pct"] < 0

        abs_ms = abs(item["target_median"] - item["baseline_median"])

        if floor is None:

            item.update(verdict="unmeasured", reason="no per-cell A/A floor for this cell")

            continue

        # The floor already takes the A/A interval's outer bound; requiring the

        # target interval to clear it as well would count the noise twice.

        if excludes_zero and abs(item["delta_pct"]) > floor and abs_ms >= min_delta_ms:

            direction = "regression" if item["delta_pct"] > 0 else "improvement"

            if rejected:

                item["verdict"] = direction

            else:

                item["verdict"] = "unconfirmed_" + direction

                item["reason"] = "not significant after Benjamini-Hochberg across the comparison family"

        else:

            item["verdict"] = "no_change"

    for key in tested:

        item = results[key]

        item["family_size"] = len(tested)

        item["mde_pct"] = minimum_detectable_effect_pct(item)

        if item["verdict"].startswith("unconfirmed_"):

            item["required_n"] = required_runs(item, len(tested), fdr)

    return results

def describe(item: dict[str, Any], label: str) -> str:

    def pct(value: float | None) -> str:

        return "?" if value is None else f"{value:+.1f}%"

    if item.get("verdict") == "no_change" and item.get("mde_pct") is not None:

        return (f"`{item['scenario']}/{item['operation']}` vs {label}: no change detected >= "

                f"{item['mde_pct']:.1f}% (observed {pct(item['delta_pct'])}, n={item['baseline_n']}/{item['target_n']})")

    return (f"`{item['scenario']}/{item['operation']}` vs {label}: {pct(item['delta_pct'])} "

            f"({item['baseline_median']:.1f} -> {item['target_median']:.1f} ms; "

            f"95% CI [{pct(item['ci_low_pct'])}, {pct(item['ci_high_pct'])}], "

            f"A/A floor {item['floor_pct']:.1f}%, p={item['p_value']:.4f}, n={item['baseline_n']}/{item['target_n']})")