"""A/A calibration: does the interval actually cover at its nominal rate?

ADR-0005 makes this a requirement rather than an assumption — "coverage must
be validated empirically." The failure being guarded against is the one the
whole project is positioned against: an interval that looks rigorous and is
quietly too narrow, so the tool confidently reports effects that aren't there.

The setup is a null world. Both arms draw from the *same* per-task success
probability, so the true Δ is exactly zero, and a well-calibrated 95% interval
should contain zero in about 95% of replications. Under-coverage here means an
unmodelled variance component, and per ADR-0005 the fix would be stratification
rather than fudging the interval wider.

These are simulations, not agent runs: no API calls, no network, no cost.
"""

from __future__ import annotations

import numpy as np
import pytest

from ablate.stats import PairedSample, bootstrap_delta

# Enough replications to distinguish "calibrated" from "3x too narrow" — the
# failure mode that matters — without making the suite slow. At 200 reps the
# standard error on a 95% coverage estimate is ~1.5pp.
N_REPLICATIONS = 200
RESAMPLES = 2_000


def _null_world(rng: np.random.Generator, n_tasks: int, n_trials: int) -> PairedSample:
    """Both arms drawn from one per-task rate, so the true effect is zero.

    Per-task rates are spread across the unit interval on purpose: tasks near
    0 and 1 are where the statistic skews and where a plain percentile interval
    under-covers, which is the case BCa's corrections exist to handle.
    """
    pairs = {}
    for i in range(n_tasks):
        p = rng.uniform(0.05, 0.95)
        pairs[f"t{i}"] = (
            list(rng.binomial(1, p, n_trials).astype(float)),
            list(rng.binomial(1, p, n_trials).astype(float)),
        )
    return PairedSample.from_pairs(pairs)


def _coverage(n_tasks: int, n_trials: int, seed: int = 20260727) -> float:
    rng = np.random.default_rng(seed)
    covered = 0
    for rep in range(N_REPLICATIONS):
        result = bootstrap_delta(
            _null_world(rng, n_tasks, n_trials), n_resamples=RESAMPLES, seed=rep
        )
        assert result.has_ci
        if result.lo <= 0.0 <= result.hi:
            covered += 1
    return covered / N_REPLICATIONS


@pytest.mark.parametrize(("n_tasks", "n_trials"), [(8, 5), (20, 5), (20, 10)])
def test_aa_coverage_is_near_nominal(n_tasks, n_trials):
    """~95% of intervals should contain zero when the truth is zero.

    The band is deliberately wide enough to absorb Monte Carlo noise at 200
    replications and the known mild conservatism of the cluster bootstrap at
    T=8, but far too tight to pass if the cluster stage were dropped — pooling
    trials would land coverage near 70%.
    """
    assert 0.90 <= _coverage(n_tasks, n_trials) <= 0.99


def _heterogeneous_world(
    rng: np.random.Generator, n_tasks: int, n_trials: int, effect_sd: float
) -> PairedSample:
    """A suite drawn from a task population whose average true effect is zero.

    Individual tasks genuinely respond to the config — some helped, some hurt,
    per-task effect ~ N(0, effect_sd) — but the population mean is zero. This
    is the realistic world, and the one that separates the two methods: the
    estimand is an average over a task *population*, so an honest interval has
    to account for which tasks landed in the suite. That variance component
    exists only at cluster level and no amount of trials per task reveals it.
    """
    pairs = {}
    for i in range(n_tasks):
        base = rng.uniform(0.25, 0.75)
        effect = rng.normal(0.0, effect_sd)
        p_a = min(max(base + effect / 2.0, 0.0), 1.0)
        p_b = min(max(base - effect / 2.0, 0.0), 1.0)
        pairs[f"t{i}"] = (
            list(rng.binomial(1, p_a, n_trials).astype(float)),
            list(rng.binomial(1, p_b, n_trials).astype(float)),
        )
    return PairedSample.from_pairs(pairs)


def test_cluster_bootstrap_covers_under_task_heterogeneity():
    """Coverage holds when tasks respond differently to the config."""
    rng = np.random.default_rng(11)
    covered = 0
    for rep in range(N_REPLICATIONS):
        result = bootstrap_delta(
            _heterogeneous_world(rng, 20, 5, effect_sd=0.15),
            n_resamples=RESAMPLES,
            seed=rep,
        )
        if result.lo <= 0.0 <= result.hi:
            covered += 1
    assert 0.90 <= covered / N_REPLICATIONS <= 0.99


def _pooled_binomial_se(a: np.ndarray, b: np.ndarray) -> float:
    """The shortcut ADR-0005 rejects: pool every trial, ignore the clusters."""
    p_a, p_b = a.mean(), b.mean()
    return (p_a * (1 - p_a) / a.size + p_b * (1 - p_b) / b.size) ** 0.5


def _sampling_sd(rng: np.random.Generator, difficulty, effect_sd, reps=1500, n_tasks=20, n_trials=5):
    """Empirical SD of Δ̂ and of a single arm's rate, plus the pooled SE."""
    deltas, arm_rates, pooled_delta_se, pooled_arm_se = [], [], [], []
    for _ in range(reps):
        base = difficulty(rng, n_tasks)
        effect = rng.normal(0.0, effect_sd, n_tasks)
        p_a = np.clip(base + effect / 2, 0, 1)
        p_b = np.clip(base - effect / 2, 0, 1)
        a = rng.binomial(1, np.repeat(p_a, n_trials).reshape(n_tasks, n_trials))
        b = rng.binomial(1, np.repeat(p_b, n_trials).reshape(n_tasks, n_trials))

        deltas.append((a.mean(axis=1) - b.mean(axis=1)).mean())
        arm_rates.append(a.mean(axis=1).mean())
        pooled_delta_se.append(_pooled_binomial_se(a, b))
        pooled_arm_se.append((a.mean() * (1 - a.mean()) / a.size) ** 0.5)

    return {
        "delta_true": float(np.std(deltas, ddof=1)),
        "delta_pooled": float(np.mean(pooled_delta_se)),
        "arm_true": float(np.std(arm_rates, ddof=1)),
        "arm_pooled": float(np.mean(pooled_arm_se)),
    }


def test_pooled_binomial_understates_a_single_arm_rate():
    """The design effect, where it is genuinely large.

    When tasks are near-deterministic — which docs/05-statistics.md says many
    are — trials within a task are strongly correlated, the effective sample
    size collapses toward T, and a pooled binomial CI on one arm's pass rate
    comes out ~2x too narrow. This is the "√N too narrow" claim in ADR-0005,
    and it holds for a *rate*.
    """
    rng = np.random.default_rng(3)
    m = _sampling_sd(rng, lambda r, t: r.choice([0.05, 0.95], t), effect_sd=0.15)
    assert m["arm_true"] / m["arm_pooled"] > 1.8


def test_pooled_binomial_understates_the_paired_delta_only_modestly():
    """Pairing does most of the work, and the report should not overclaim.

    ADR-0005's headline magnitude ("2–3x too narrow") is measured on a single
    arm's rate. For the *paired difference* — which is what `ablate compare`
    actually reports — pairing cancels between-task difficulty, and the pooled
    shortcut is only mildly too narrow: ~1.05x at moderate effect heterogeneity,
    ~1.3x when tasks respond very differently to the config.

    That is still a real error and still worth the cluster bootstrap, because
    the residual gap is exactly the between-task effect variance and no number
    of trials reveals it. But the honest figure for the paired delta is "modest
    and heterogeneity-dependent," not "3x", and the README's method section
    should say so rather than inherit the single-arm number.
    """
    mid = lambda r, t: r.uniform(0.25, 0.75, t)  # noqa: E731

    rng = np.random.default_rng(3)
    moderate = _sampling_sd(rng, mid, effect_sd=0.15)
    assert 1.0 < moderate["delta_true"] / moderate["delta_pooled"] < 1.15

    rng = np.random.default_rng(3)
    severe = _sampling_sd(rng, mid, effect_sd=0.30)
    assert 1.2 < severe["delta_true"] / severe["delta_pooled"] < 1.5

    # The gap grows with heterogeneity — that is the component pooling misses.
    assert (
        severe["delta_true"] / severe["delta_pooled"]
        > moderate["delta_true"] / moderate["delta_pooled"]
    )


def test_pooled_binomial_is_too_wide_for_deterministic_tasks():
    """The regime where the shortcut errs the other way.

    With near-deterministic tasks and a small true effect, the paired
    difference has very little variance, but the pooled estimator still carries
    the full between-task difficulty spread — so its interval is roughly 2x too
    *wide*. A tool that hardcoded a "pooling is too narrow" correction factor
    would be wrong here. The cluster bootstrap needs no such factor: it adapts
    to whichever regime the suite is actually in, which is the real argument
    for it.
    """
    rng = np.random.default_rng(3)
    m = _sampling_sd(rng, lambda r, t: r.choice([0.05, 0.95], t), effect_sd=0.15)
    assert m["delta_true"] / m["delta_pooled"] < 0.75
