"""Tests for the inferential engine.

The bootstrap is the one place in this codebase where a bug is a silent wrong
answer rather than a crash, so these tests check it against constructed cases
with known answers rather than against its own output.
"""

from __future__ import annotations

import numpy as np
import pytest

from ablate.stats import (
    MIN_TASKS_FOR_CI,
    PairedSample,
    bootstrap_delta,
    mde_continuous,
    mde_pass_rate,
    required_design,
)


def _sample(n_tasks: int, a_rate: float, b_rate: float, n_trials: int = 5) -> PairedSample:
    """A deterministic sample where each task hits its rate exactly."""
    def outcomes(rate: float) -> list[float]:
        k = round(rate * n_trials)
        return [1.0] * k + [0.0] * (n_trials - k)

    return PairedSample.from_pairs(
        {f"task-{i:02d}": (outcomes(a_rate), outcomes(b_rate)) for i in range(n_tasks)}
    )


# --- point estimate ---------------------------------------------------------


def test_delta_is_the_mean_of_per_task_differences():
    sample = PairedSample.from_pairs(
        {
            "t1": ([1, 1, 1, 1, 0], [1, 1, 0, 0, 0]),  # 0.8 - 0.4 = +0.4
            "t2": ([0, 0, 0, 0, 0], [1, 0, 0, 0, 0]),  # 0.0 - 0.2 = -0.2
        }
    )
    result = bootstrap_delta(sample, seed=1)
    assert result.delta == pytest.approx(0.1)


def test_pairing_is_by_task_not_by_pooled_trial():
    """A task with more trials must not dominate the estimate.

    Pooling trials would weight `heavy` 10x; pairing by task weights it 1x.
    This is the error ADR-0005 exists to reject, so it gets a test.
    """
    sample = PairedSample.from_pairs(
        {
            "heavy": ([1.0] * 50, [0.0] * 50),  # d = +1.0, 50 trials
            "light": ([0.0] * 5, [1.0] * 5),  # d = -1.0, 5 trials
        }
    )
    assert bootstrap_delta(sample, seed=1).delta == pytest.approx(0.0)


def test_weights_shift_the_estimate():
    pairs = {"t1": ([1, 1], [0, 0]), "t2": ([0, 0], [1, 1])}
    weighted = PairedSample.from_pairs(pairs, weights={"t1": 3.0, "t2": 1.0})
    # (3*(+1) + 1*(-1)) / 4
    assert bootstrap_delta(weighted, seed=1).delta == pytest.approx(0.5)


def test_tasks_missing_an_arm_are_dropped():
    sample = PairedSample.from_pairs({"paired": ([1], [0]), "orphan": ([1], [])})
    assert sample.task_ids == ("paired",)


# --- interval ---------------------------------------------------------------


def test_no_interval_below_the_cluster_floor():
    result = bootstrap_delta(_sample(MIN_TASKS_FOR_CI - 1, 0.6, 0.4), seed=1)
    assert not result.has_ci
    assert result.ci_withheld is not None
    assert str(MIN_TASKS_FOR_CI) in result.ci_withheld
    assert result.delta == pytest.approx(0.2)  # the point estimate still stands


def test_interval_emitted_at_the_floor():
    assert bootstrap_delta(_sample(MIN_TASKS_FOR_CI, 0.6, 0.4), seed=1).has_ci


def test_interval_brackets_the_point_estimate():
    result = bootstrap_delta(_sample(20, 0.6, 0.4), seed=1)
    assert result.lo <= result.delta <= result.hi


def test_zero_variance_sample_gives_a_degenerate_interval():
    """Every task identical: there is nothing to resample, so the interval
    collapses onto the point estimate rather than inventing width."""
    result = bootstrap_delta(_sample(20, 1.0, 0.0), seed=1)
    assert result.delta == pytest.approx(1.0)
    assert result.lo == pytest.approx(1.0)
    assert result.hi == pytest.approx(1.0)


def test_interval_narrows_as_tasks_are_added():
    rng = np.random.default_rng(0)

    def noisy(n_tasks: int) -> PairedSample:
        r = np.random.default_rng(7)
        return PairedSample.from_pairs(
            {
                f"t{i}": (
                    list(r.binomial(1, 0.6, 5).astype(float)),
                    list(r.binomial(1, 0.5, 5).astype(float)),
                )
                for i in range(n_tasks)
            }
        )

    del rng
    narrow = bootstrap_delta(noisy(60), seed=3)
    wide = bootstrap_delta(noisy(10), seed=3)
    assert (narrow.hi - narrow.lo) < (wide.hi - wide.lo)


def test_seed_makes_the_interval_reproducible():
    sample = _sample(12, 0.6, 0.4)
    first = bootstrap_delta(sample, seed=42)
    second = bootstrap_delta(sample, seed=42)
    assert (first.lo, first.hi) == (second.lo, second.hi)


def test_unbalanced_design_is_analyzable():
    """A run stopped by the budget ceiling leaves ragged arms. Because the
    queue is shuffled, that partial data is still usable — so it must not
    crash, and each task must resample to its own N."""
    pairs = {
        f"t{i}": ([1.0] * (3 + i % 3), [0.0] * (5 - i % 4)) for i in range(12)
    }
    result = bootstrap_delta(PairedSample.from_pairs(pairs), seed=1)
    assert result.has_ci
    assert result.delta == pytest.approx(1.0)


# --- relative mode ----------------------------------------------------------


def test_relative_delta_is_a_fraction_of_arm_b():
    pairs = {f"t{i}": ([80.0] * 5, [100.0] * 5) for i in range(10)}
    result = bootstrap_delta(PairedSample.from_pairs(pairs), relative=True, seed=1)
    assert result.delta == pytest.approx(-0.2)
    assert result.relative is True


def test_relative_delta_survives_a_zero_baseline_task():
    pairs = {f"t{i}": ([1.0] * 5, [0.0] * 5) for i in range(10)}
    result = bootstrap_delta(PairedSample.from_pairs(pairs), relative=True, seed=1)
    assert not result.has_ci  # division by a zero baseline is undefined, not zero
    assert result.ci_withheld is not None


# --- power ------------------------------------------------------------------


@pytest.mark.parametrize(
    ("n_tasks", "n_trials", "expected_se"),
    [
        (8, 3, 0.145),
        (8, 5, 0.113),
        (20, 5, 0.072),
        (20, 10, 0.051),
        (40, 10, 0.036),
        (100, 10, 0.023),
    ],
)
def test_standard_error_matches_the_documented_table(n_tasks, n_trials, expected_se):
    """docs/05-statistics.md's MDE table, reproduced from the formula.

    Pinning the published table to the implementation keeps the two from
    drifting apart — which they had: the T=8/N=3 row read 0.132, which the
    document's own formula does not produce. The formula was taken as
    authoritative and the doc corrected to 0.145.
    """
    from ablate.stats import standard_error

    se = standard_error(n_tasks, n_trials, within_task_variance=0.5)
    assert se == pytest.approx(expected_se, abs=0.0005)


@pytest.mark.parametrize(
    ("n_tasks", "n_trials", "expected_mde"),
    [
        (8, 3, 0.41),
        (8, 5, 0.32),
        (20, 5, 0.20),
        (20, 10, 0.14),
        (40, 10, 0.10),
        (100, 10, 0.064),
    ],
)
def test_pass_rate_mde_matches_the_documented_table(n_tasks, n_trials, expected_mde):
    assert mde_pass_rate(n_tasks, n_trials) == pytest.approx(expected_mde, abs=0.005)


@pytest.mark.parametrize(
    ("n_tasks", "n_trials", "expected"), [(8, 5, 0.16), (20, 5, 0.10), (20, 10, 0.07)]
)
def test_token_mde_matches_the_documented_table(n_tasks, n_trials, expected):
    assert mde_continuous(n_tasks, n_trials) == pytest.approx(expected, abs=0.005)


def test_token_mde_beats_pass_rate_mde_at_the_same_design():
    """ADR-0010's entire argument: at an affordable N the token metric is
    detectable and the pass-rate metric is not."""
    assert mde_continuous(20, 5) < mde_pass_rate(20, 5) / 2


def test_mde_improves_at_lopsided_pass_rates():
    assert mde_pass_rate(20, 5, p_a=0.8, p_b=0.8) < mde_pass_rate(20, 5)


def test_required_design_achieves_its_target():
    req = required_design(0.05, n_trials=10)
    assert req.achieved_mde <= 0.05
    assert req.n_runs == req.n_tasks * req.n_trials * 2
    assert req.feasible


def test_required_design_needs_more_tasks_for_a_smaller_effect():
    assert required_design(0.03, n_trials=10).n_tasks > required_design(0.10, n_trials=10).n_tasks


def test_required_design_reports_infeasible_targets():
    """Below the between-task heterogeneity floor, no affordable T resolves it.
    The tool must say "not at any price you'll pay" rather than print a number."""
    assert not required_design(0.001, n_trials=10, max_tasks=1000).feasible


def test_required_design_rejects_a_nonpositive_target():
    with pytest.raises(ValueError, match="positive"):
        required_design(0.0)


def test_a_metric_no_task_reported_yields_no_estimate():
    """An agent that does not report token usage produces no token sample at
    all. The delta is NaN rather than 0 — there is nothing to average, and 0
    would be a measurement."""
    import math

    result = bootstrap_delta(PairedSample.from_pairs({}), seed=1)
    assert math.isnan(result.delta)
    assert not result.has_ci
    assert "in both arms" in result.ci_withheld
