"""Paired cluster bootstrap with BCa intervals (docs/05-statistics.md, ADR-0005).

The estimand, precisely:

    Δ = the average, over the tasks in this suite, of the per-task difference
        in metric between config A and config B, for this agent at this model
        version.

Pairing is at **task** level, not trial level: current frontier models expose
no seed or temperature control, so trial *r* of arm A and trial *r* of arm B
share no randomness and common-random-numbers variance reduction is simply
unavailable. Pairing by task still cancels between-task difficulty variance,
which is the largest variance component by far.

Resampling is therefore two-stage, and the outer stage is the one that
matters: resample **tasks** with replacement (the clusters), then resample
trials within each arm of each drawn task. Skipping the cluster stage — i.e.
pooling all trials and computing a binomial CI — treats correlated trials as
independent observations and yields an interval roughly √N too narrow. That
error is the reason this tool exists; see ADR-0005.

This module is deliberately ignorant of `TrialRecord`, tasks, and configs. It
takes arrays of numbers so it can be tested against known answers.
"""

from __future__ import annotations

from dataclasses import dataclass
from statistics import NormalDist

import numpy as np

#: Below this many tasks, no interval is emitted at all. Bootstrap coverage at
#: five clusters is not honest, and printing a wide interval implies more rigor
#: than exists. `ablate try` (3 tasks) therefore cannot produce a CI, which is
#: correct and which its report says out loud.
MIN_TASKS_FOR_CI = 8

DEFAULT_RESAMPLES = 10_000

_NORM = NormalDist()


@dataclass(frozen=True)
class PairedSample:
    """Per-task trial outcomes for one metric, in both arms.

    `a[i]` and `b[i]` are the raw per-trial values for task `task_ids[i]` under
    config A and config B. They need not be the same length — a run stopped by
    the budget ceiling leaves an unbalanced design, and because the work queue
    is shuffled that partial data is still analyzable.
    """

    task_ids: tuple[str, ...]
    a: tuple[np.ndarray, ...]
    b: tuple[np.ndarray, ...]
    weights: np.ndarray

    @classmethod
    def from_pairs(
        cls,
        pairs: dict[str, tuple[list[float], list[float]]],
        weights: dict[str, float] | None = None,
    ) -> PairedSample:
        """Build from `{task_id: (arm_a_values, arm_b_values)}`.

        Tasks with no trials in either arm are dropped: they carry no
        information and would make the per-task mean undefined.
        """
        task_ids = tuple(sorted(t for t, (va, vb) in pairs.items() if va and vb))
        return cls(
            task_ids=task_ids,
            a=tuple(np.asarray(pairs[t][0], dtype=float) for t in task_ids),
            b=tuple(np.asarray(pairs[t][1], dtype=float) for t in task_ids),
            weights=np.array(
                [1.0 if weights is None else weights.get(t, 1.0) for t in task_ids],
                dtype=float,
            ),
        )

    @property
    def n_tasks(self) -> int:
        return len(self.task_ids)

    def task_means(self) -> tuple[np.ndarray, np.ndarray]:
        """Observed per-task means in each arm."""
        return (
            np.array([x.mean() for x in self.a]),
            np.array([x.mean() for x in self.b]),
        )


@dataclass(frozen=True)
class BootstrapResult:
    """A point estimate and, when the design earns one, an interval.

    `lo`/`hi` are `None` exactly when `ci_withheld` explains why. Callers must
    render the withholding reason rather than substituting a bare point
    estimate — a number with a silently missing interval is the failure mode
    this whole module is built to prevent.
    """

    delta: float
    lo: float | None
    hi: float | None
    alpha: float
    n_tasks: int
    n_resamples: int
    relative: bool
    method: str
    ci_withheld: str | None = None

    @property
    def has_ci(self) -> bool:
        return self.lo is not None and self.hi is not None


def _weighted(values: np.ndarray, w: np.ndarray, axis: int = -1) -> np.ndarray:
    return (values * w).sum(axis=axis) / w.sum(axis=axis)


def _statistic(
    mean_a: np.ndarray, mean_b: np.ndarray, w: np.ndarray, relative: bool, axis: int = -1
) -> np.ndarray:
    """Δ̂ from per-task arm means.

    Absolute: the weighted mean of the per-task differences. Relative: that
    same numerator over the weighted mean of arm B, so the ratio is formed
    inside each resample rather than after the fact — a ratio of averages,
    bootstrapped as a unit, which is what makes its interval honest.
    """
    delta = np.asarray(_weighted(mean_a - mean_b, w, axis=axis), dtype=float)
    if not relative:
        return delta
    base = np.asarray(_weighted(mean_b, w, axis=axis), dtype=float)
    # A task suite whose arm-B mean is zero has no baseline to be relative to.
    # NaN propagates to `ci_withheld` rather than silently becoming 0 or inf.
    out = np.full(delta.shape, np.nan, dtype=float)
    np.divide(delta, base, out=out, where=base != 0)
    return out


def _resample_arm(
    rng: np.random.Generator, arrays: tuple[np.ndarray, ...], task_idx: np.ndarray
) -> np.ndarray:
    """Trial-level resampling, vectorized over (resample, task-slot).

    Tasks are grouped by trial count so each group can be gathered as one
    rectangular array. Ragged designs (a run cut short by `--max-spend`) are
    therefore handled without falling back to a Python loop over resamples,
    and every task is resampled to exactly its own N — never padded up to the
    maximum, which would quietly shrink its variance.
    """
    sizes = np.array([len(x) for x in arrays])
    out = np.empty(task_idx.shape, dtype=float)

    for size in np.unique(sizes):
        members = np.flatnonzero(sizes == size)
        lookup = np.full(len(arrays), -1, dtype=int)
        lookup[members] = np.arange(len(members))

        selected = lookup[task_idx]
        mask = selected >= 0
        k = int(mask.sum())
        if k == 0:
            continue

        block = np.stack([arrays[i] for i in members])
        draws = rng.integers(0, size, size=(k, size))
        out[mask] = block[selected[mask][:, None], draws].mean(axis=1)

    return out


@dataclass(frozen=True)
class BCaCorrections:
    """BCa's two correction terms.

    `bias` (z0) corrects median bias. `acceleration` corrects for the
    statistic's variance changing with its own value, which it does whenever
    per-task rates sit near 0 or 1 — exactly where a lot of tasks sit, and
    where a plain percentile interval under-covers.

    `applied=False` means the point estimate fell outside the resample
    distribution entirely, so the bias correction is undefined. The interval
    degrades to a plain percentile one and callers label it as such, rather
    than emitting a silently wrong BCa.
    """

    bias: float
    acceleration: float
    applied: bool


def bca_corrections(
    resampled: np.ndarray, point: float, jackknife: np.ndarray
) -> BCaCorrections:
    """Compute z0 and the acceleration from a resample distribution.

    `jackknife` holds the statistic recomputed with each **cluster** deleted in
    turn — a task, not a trial. The acceleration has to describe how the
    statistic responds to the resampling unit, and the resampling unit is the
    task.
    """
    below = float((resampled < point).mean())
    if not 0.0 < below < 1.0:
        return BCaCorrections(0.0, 0.0, applied=False)

    acceleration = 0.0
    if jackknife.size and np.all(np.isfinite(jackknife)):
        centered = jackknife.mean() - jackknife
        denom = 6.0 * float((centered**2).sum()) ** 1.5
        if denom:
            acceleration = float((centered**3).sum() / denom)

    return BCaCorrections(_NORM.inv_cdf(below), acceleration, applied=True)


def bca_bounds(
    resampled: np.ndarray, corrections: BCaCorrections, alpha: float
) -> tuple[float, float]:
    """Adjusted percentile bounds from a resample distribution."""

    def endpoint(z: float) -> float:
        adjusted = corrections.bias + z
        divisor = 1.0 - corrections.acceleration * adjusted
        shifted = corrections.bias + (adjusted / divisor if divisor else adjusted)
        return min(max(_NORM.cdf(shifted), 0.0), 1.0)

    lo_q = endpoint(_NORM.inv_cdf(alpha / 2.0))
    hi_q = endpoint(_NORM.inv_cdf(1.0 - alpha / 2.0))
    lo, hi = np.percentile(resampled, [100.0 * lo_q, 100.0 * hi_q])
    return float(min(lo, hi)), float(max(lo, hi))


def _jackknife_values(sample: PairedSample, relative: bool) -> np.ndarray:
    """The statistic with each task deleted in turn."""
    mean_a, mean_b = sample.task_means()
    w = sample.weights
    return np.array(
        [
            _statistic(
                np.delete(mean_a, i), np.delete(mean_b, i), np.delete(w, i), relative, axis=0
            )
            for i in range(sample.n_tasks)
        ],
        dtype=float,
    )


def bootstrap_delta(
    sample: PairedSample,
    *,
    relative: bool = False,
    n_resamples: int = DEFAULT_RESAMPLES,
    alpha: float = 0.05,
    seed: int | None = None,
) -> BootstrapResult:
    """Δ̂ with a BCa paired-cluster-bootstrap interval.

    `relative=True` reports the delta as a fraction of arm B — the form the
    token metric is quoted in ("−29.4% [−31.1, −27.8]").

    `seed` fixes the resampling only. It cannot make the *agent* reproducible;
    nothing can. Reports say so.
    """
    n = sample.n_tasks
    if n == 0:
        # No task had usable values in both arms. Common and legitimate: an
        # agent that does not report token usage produces no token metric at
        # all (see the Codex adapter and Q1). The point estimate is NaN rather
        # than 0 — there is nothing to average, and 0 would be a measurement.
        return BootstrapResult(
            delta=float("nan"),
            lo=None,
            hi=None,
            alpha=alpha,
            n_tasks=0,
            n_resamples=0,
            relative=relative,
            method="none",
            ci_withheld="no task reported this metric in both arms",
        )

    mean_a, mean_b = sample.task_means()
    point = float(_statistic(mean_a, mean_b, sample.weights, relative, axis=0))

    if n < MIN_TASKS_FOR_CI:
        return BootstrapResult(
            delta=point,
            lo=None,
            hi=None,
            alpha=alpha,
            n_tasks=n,
            n_resamples=0,
            relative=relative,
            method="none",
            ci_withheld=(
                f"{n} task{'s' if n != 1 else ''} — an interval needs at least "
                f"{MIN_TASKS_FOR_CI}. Bootstrap coverage at this few clusters is not "
                f"trustworthy, so no interval is reported."
            ),
        )

    rng = np.random.default_rng(seed)
    task_idx = rng.integers(0, n, size=(n_resamples, n))

    resampled = _statistic(
        _resample_arm(rng, sample.a, task_idx),
        _resample_arm(rng, sample.b, task_idx),
        sample.weights[task_idx],
        relative,
        axis=1,
    )
    resampled = resampled[np.isfinite(resampled)]
    if resampled.size == 0:
        return BootstrapResult(
            delta=point,
            lo=None,
            hi=None,
            alpha=alpha,
            n_tasks=n,
            n_resamples=n_resamples,
            relative=relative,
            method="none",
            ci_withheld="every resample was degenerate — the metric has no variation to resample",
        )

    corrected = bca_corrections(resampled, point, _jackknife_values(sample, relative))
    method = (
        "paired cluster bootstrap, BCa"
        if corrected.applied
        else "paired cluster bootstrap, percentile (BCa correction undefined)"
    )
    lo, hi = bca_bounds(resampled, corrected, alpha)

    return BootstrapResult(
        delta=point,
        lo=float(min(lo, hi)),
        hi=float(max(lo, hi)),
        alpha=alpha,
        n_tasks=n,
        n_resamples=n_resamples,
        relative=relative,
        method=method,
        ci_withheld=None,
    )
