"""Tests for the screening regression.

These are mostly recovery tests: given a world with known section effects, does
the fit rank them correctly and does it say "nothing here" about the sections
that genuinely have no effect. Ranking is what stage 1 is *for*, so that is
what gets pinned — not the coefficient magnitudes, which ridge deliberately
shrinks.
"""

from __future__ import annotations

import numpy as np
import pytest

from ablate.ablation import build_design
from ablate.ablation.regression import LAMBDA_GRID, fit_sections
from ablate.errors import AblateError

K = 8
TRUE_BETA = np.array([0.30, 0.15, 0.0, 0.0, -0.20, 0.0, 0.05, 0.0])
REAL_EFFECTS = np.abs(TRUE_BETA) > 0.1


def _simulate(n_tasks: int = 12, seed: int = 0, true_beta: np.ndarray = TRUE_BETA):
    """A world where keeping section k shifts the pass probability by beta_k,
    on top of a task difficulty that varies widely."""
    design = build_design([f"s{i + 1:02d}" for i in range(len(true_beta))], seed=seed)
    matrix = design.matrix()
    rng = np.random.default_rng(1000 + seed)
    difficulty = rng.uniform(0.2, 0.8, n_tasks)

    masks, tasks, y = [], [], []
    for t in range(n_tasks):
        for v in range(design.n_variants):
            p = np.clip(
                difficulty[t] + matrix[v] @ true_beta - true_beta.sum() / 2, 0.02, 0.98
            )
            masks.append(matrix[v])
            tasks.append(f"task{t}")
            y.append(float(rng.random() < p))

    return design, np.array(masks), np.array(tasks), np.array(y)


def _fit(n_tasks: int = 12, seed: int = 0, **kw):
    design, masks, tasks, y = _simulate(n_tasks, seed)
    return fit_sections(design.section_ids, masks, tasks, y, seed=seed, **kw)


# --- recovery ---------------------------------------------------------------


def test_ranking_recovers_the_sections_that_matter():
    """Stage 1's job is a ranked shortlist, and it is judged on that."""
    hits = 0
    truth = set(np.argsort(-np.abs(TRUE_BETA))[:3])
    for seed in range(10):
        fit = _fit(seed=seed, n_resamples=120)
        beta = np.array([e.beta for e in fit.effects])
        hits += len(truth & set(np.argsort(-np.abs(beta))[:3]))
    assert hits / 30 >= 0.6  # measured ~0.8; the floor guards against regressions


def test_the_largest_true_effect_is_found():
    fit = _fit(seed=3, n_resamples=200)
    beta = np.array([e.beta for e in fit.effects])
    assert int(np.argmax(np.abs(beta))) == int(np.argmax(np.abs(TRUE_BETA)))


def test_signs_of_real_effects_are_recovered():
    fit = _fit(seed=5, n_resamples=200)
    beta = np.array([e.beta for e in fit.effects])
    assert np.all(np.sign(beta[REAL_EFFECTS]) == np.sign(TRUE_BETA[REAL_EFFECTS]))


def test_sections_with_no_effect_have_intervals_spanning_zero():
    """The deletion candidates. Getting this wrong recommends keeping dead
    weight, or worse, recommends deleting something that matters.

    Measured coverage here is ~0.88 against a nominal 0.95 -- in a world where
    *other* sections have real effects. In a purely null world it measures
    ~0.97. The residual gap is ridge bias at small T with correlated columns,
    not the penalty-selection problem, which `test_retuning_lambda_widens_the_interval`
    covers.

    Screening intervals are therefore approximately calibrated, not exactly,
    and that is the concrete reason stage 2 is not optional. It is disclosed in
    the report rather than papered over -- unlike the paired-comparison
    interval in `ablate compare`, which is calibrated to nominal and tested as
    such in test_calibration.py.
    """
    spans = []
    for seed in range(20):
        fit = _fit(seed=seed, n_resamples=300)
        spans += [e.spans_zero for k, e in enumerate(fit.effects) if TRUE_BETA[k] == 0]
    assert 0.78 <= np.mean(spans) <= 1.0  # measured 0.88 at 40 seeds


def test_retuning_lambda_widens_the_interval():
    """lambda is chosen from the data, so holding it fixed across resamples
    treats a fitted quantity as a known constant.

    How much that costs depends on the world, and the difference is large
    enough to be worth stating both ways. In a purely null world the fixed-lambda
    interval is ~4x too narrow (89% coverage vs 97% re-tuned). In this world,
    where other sections carry real effects, the gap is ~1.1x (82% vs 88%).

    The null world is the decision-relevant one -- a deletion candidate *is* a
    section whose effect looks like zero -- so the expensive path is the
    default. The wide interval is the honest one; the narrow one merely looked
    authoritative.
    """
    ratios = []
    for seed in range(8):
        retuned = _fit(seed=seed, n_resamples=300)
        fixed = _fit(seed=seed, n_resamples=300, retune_lambda=False)
        ratios.append(
            np.mean([e.hi - e.lo for e in retuned.effects])
            / np.mean([e.hi - e.lo for e in fixed.effects])
        )
    assert np.mean(ratios) > 1.05  # measured ~1.14 here, ~4x in a null world


def test_retuning_is_the_default():
    """The cheap option is the wrong one, so it must not be reachable by
    accident."""
    import inspect

    assert inspect.signature(fit_sections).parameters["retune_lambda"].default is True


def test_a_world_with_no_effects_at_all_produces_no_findings():
    design, masks, tasks, y = _simulate(seed=2, true_beta=np.zeros(K))
    fit = fit_sections(design.section_ids, masks, tasks, y, seed=2, n_resamples=200)
    assert sum(not e.spans_zero for e in fit.effects) <= 1  # ~alpha * K


# --- ridge behaviour --------------------------------------------------------


def test_coefficients_are_shrunk_toward_zero():
    """Ridge is a biased estimator and the report has to say so. Intervals
    cover the *penalized* estimand, not the true effect, which is why stage 1
    is labelled screening and stage 2 is not optional."""
    ratios = []
    for seed in range(6):
        fit = _fit(seed=seed, n_resamples=80)
        beta = np.array([e.beta for e in fit.effects])
        ratios.append(np.mean(np.abs(beta[REAL_EFFECTS]) / np.abs(TRUE_BETA[REAL_EFFECTS])))
    mean_ratio = float(np.mean(ratios))
    assert 0.4 < mean_ratio < 1.0, "should shrink, but not annihilate"


def test_no_coefficient_is_forced_exactly_to_zero():
    """The LASSO failure ADR-0007 rejects: zeroing a coefficient destroys the
    distinction between "measured to be zero" and "shrunk to zero", and that
    distinction is the product."""
    fit = _fit(seed=1, n_resamples=80)
    assert all(e.beta != 0.0 for e in fit.effects)


def test_a_larger_penalty_shrinks_further():
    small = _fit(seed=1, lam=0.1, n_resamples=60)
    large = _fit(seed=1, lam=100.0, n_resamples=60)
    assert np.abs([e.beta for e in large.effects]).sum() < np.abs(
        [e.beta for e in small.effects]
    ).sum()


def test_penalty_is_chosen_from_the_grid():
    assert _fit(seed=1, n_resamples=60).lam in set(LAMBDA_GRID)


def test_explicit_penalty_overrides_cross_validation():
    assert _fit(seed=1, lam=2.5, n_resamples=60).lam == 2.5


# --- task fixed effects -----------------------------------------------------


def test_task_difficulty_does_not_leak_into_the_coefficients():
    """The within transformation must absorb alpha_t exactly. If it did not,
    a suite with wildly varying task difficulty would produce spurious betas."""
    design = build_design([f"s{i + 1:02d}" for i in range(4)], seed=1)
    matrix = design.matrix()

    masks, tasks, y = [], [], []
    for t, difficulty in enumerate([0.05, 0.35, 0.65, 0.95] * 3):
        for v in range(design.n_variants):
            masks.append(matrix[v])
            tasks.append(f"task{t}")
            y.append(difficulty)  # no section effect whatsoever

    fit = fit_sections(design.section_ids, np.array(masks), np.array(tasks), np.array(y), seed=1)
    assert np.allclose([e.beta for e in fit.effects], 0.0, atol=1e-9)


def test_a_constant_outcome_yields_zero_effects():
    design, masks, tasks, _ = _simulate(seed=1)
    fit = fit_sections(design.section_ids, masks, tasks, np.ones(len(masks)), seed=1)
    assert np.allclose([e.beta for e in fit.effects], 0.0, atol=1e-9)


# --- intervals --------------------------------------------------------------


def test_intervals_bracket_their_point_estimate():
    fit = _fit(seed=4, n_resamples=150)
    assert all(e.lo <= e.beta <= e.hi for e in fit.effects)


def test_intervals_are_withheld_below_the_cluster_floor():
    """Same 8-task floor as everywhere else — screening does not get a
    looser standard just because it is cheaper."""
    fit = _fit(n_tasks=5, seed=1)
    assert fit.ci_withheld is not None
    assert all(not e.has_ci for e in fit.effects)
    assert all(e.beta != 0.0 for e in fit.effects)  # point estimates still stand


def test_sections_without_an_interval_count_as_spanning_zero():
    """No interval means no evidence of an effect, which must not read as
    evidence of no effect — but it also cannot claim a finding."""
    fit = _fit(n_tasks=5, seed=1)
    assert all(e.spans_zero for e in fit.effects)


def test_seed_makes_intervals_reproducible():
    first, second = _fit(seed=9, n_resamples=100), _fit(seed=9, n_resamples=100)
    assert [e.lo for e in first.effects] == [e.lo for e in second.effects]


def test_more_tasks_tighten_the_intervals():
    narrow = _fit(n_tasks=30, seed=2, n_resamples=100)
    wide = _fit(n_tasks=9, seed=2, n_resamples=100)
    assert np.mean([e.hi - e.lo for e in narrow.effects]) < np.mean(
        [e.hi - e.lo for e in wide.effects]
    )


def test_p_values_are_bounded():
    fit = _fit(seed=1, n_resamples=100)
    assert all(0.0 <= e.p_value <= 1.0 for e in fit.effects)


def test_a_strong_effect_gets_a_small_p_value():
    fit = _fit(seed=3, n_resamples=200)
    assert fit.by_id("s01").p_value < 0.05


# --- input validation -------------------------------------------------------


def test_mismatched_design_width_is_rejected():
    with pytest.raises(AblateError, match="design matrix"):
        fit_sections(["s01", "s02"], np.zeros((10, 3)), np.array(["t"] * 10), np.zeros(10))


def test_mismatched_observation_counts_are_rejected():
    with pytest.raises(AblateError, match="same observations"):
        fit_sections(["s01"], np.zeros((10, 1)), np.array(["t"] * 9), np.zeros(10))


def test_empty_input_is_rejected():
    with pytest.raises(AblateError, match="no observations"):
        fit_sections(["s01"], np.zeros((0, 1)), np.array([]), np.array([]))
