"""Tests for the screening design generator.

ADR-0007: "a collinear design produces confidently wrong βs." Conditioning is
therefore the thing under test — not that a design is produced, but that a bad
one never is.
"""

from __future__ import annotations

import numpy as np
import pytest

from ablate.ablation import EXTRA_VARIANTS, MAX_PAIRWISE_CORR, build_design
from ablate.ablation.design import _conditioning
from ablate.errors import AblateError


def _ids(k: int) -> list[str]:
    return [f"s{i + 1:02d}" for i in range(k)]


# --- shape ------------------------------------------------------------------


def test_design_has_2k_plus_8_variants():
    design = build_design(_ids(12), seed=1)
    assert design.n_variants == 2 * 12 + EXTRA_VARIANTS == 32


def test_first_variant_keeps_everything():
    """The reference: the config exactly as the user wrote it."""
    design = build_design(_ids(6), seed=1)
    assert all(design.masks[design.reference_index])


def test_second_variant_drops_everything():
    """The empty-config floor — the question most users actually have,
    answered for free by the design."""
    design = build_design(_ids(6), seed=1)
    assert not any(design.masks[design.floor_index])


def test_random_block_is_neither_all_kept_nor_all_dropped_on_average():
    design = build_design(_ids(10), seed=3)
    kept = design.matrix()[2:].mean()
    assert 0.3 < kept < 0.7  # p = 0.5, with room for sampling noise


def test_trial_count_is_one_per_variant_task_cell():
    design = build_design(_ids(12), seed=1)
    assert design.n_trials(8) == 32 * 8 == 256  # docs/06-cost.md's worked example


def test_seed_makes_the_design_reproducible():
    assert build_design(_ids(8), seed=42).masks == build_design(_ids(8), seed=42).masks


def test_different_seeds_give_different_designs():
    assert build_design(_ids(8), seed=1).masks != build_design(_ids(8), seed=2).masks


# --- conditioning -----------------------------------------------------------


@pytest.mark.parametrize("k", range(1, 25))
def test_every_supported_k_yields_a_conditioned_design(k):
    """K ranges from 1 to MAX_SECTIONS; all of it has to work."""
    design = build_design(_ids(k), seed=k)
    assert design.conditioning.ok
    assert not design.conditioning.constant_columns


def test_no_section_is_constant_across_variants():
    """A section that never varies has no estimable effect."""
    matrix = build_design(_ids(12), seed=5).matrix()
    assert matrix.std(axis=0).min() > 0


def test_collinearity_stays_under_the_threshold():
    design = build_design(_ids(12), seed=5)
    assert design.conditioning.max_abs_corr < MAX_PAIRWISE_CORR


def test_conditioning_flags_a_constant_column():
    matrix = np.array([[1.0, 1.0], [0.0, 1.0], [1.0, 1.0]])
    result = _conditioning(matrix)
    assert not result.ok
    assert result.constant_columns == (1,)
    assert "never varied" in result.describe(("s01", "s02"))


def test_conditioning_flags_a_collinear_pair():
    """Two sections that always move together cannot be told apart."""
    matrix = np.array([[1.0, 1.0], [0.0, 0.0], [1.0, 1.0], [0.0, 0.0]])
    result = _conditioning(matrix)
    assert not result.ok
    assert result.worst_pair is not None
    assert "collinear" in result.describe(("s01", "s02"))


def test_conditioning_reports_constant_columns_before_collinearity():
    """A constant column has undefined correlation, so calling it collinear
    would misdescribe the problem."""
    matrix = np.array([[1.0, 1.0], [1.0, 0.0], [1.0, 1.0]])
    result = _conditioning(matrix)
    assert result.constant_columns == (0,)
    assert result.worst_pair is None


def test_single_section_design_is_trivially_conditioned():
    result = _conditioning(np.array([[1.0], [0.0], [1.0]]))
    assert result.ok
    assert result.max_abs_corr == 0.0


# --- refusals ---------------------------------------------------------------


def test_refuses_when_every_section_is_pinned():
    with pytest.raises(AblateError, match="nothing to ablate"):
        build_design([])


def test_refuses_a_design_too_short_to_condition():
    with pytest.raises(AblateError, match="at least 3 variants"):
        build_design(_ids(4), n_variants=2)


def test_gives_up_loudly_rather_than_fitting_a_bad_design():
    """Silently proceeding would produce a confident ranking that means
    nothing, and no caveat repairs that."""
    with pytest.raises(AblateError, match="well-conditioned"):
        build_design(_ids(4), n_variants=3, max_attempts=50)


def test_failure_message_suggests_a_way_forward():
    with pytest.raises(AblateError, match="granularity bullet"):
        build_design(_ids(4), n_variants=3, max_attempts=50)
