"""Tests for evals.stats — statistical analysis module."""

from __future__ import annotations

import math

from evals.stats import (
    BinaryPair,
    analyze_condition_matrix,
    analyze_paired_results,
    any_success_aggregation,
    bootstrap_ci,
    cohens_h,
    holm_bonferroni,
    majority_vote_aggregation,
    mcnemar_test,
    power_report,
)

# ---------------------------------------------------------------------------
# McNemar's test
# ---------------------------------------------------------------------------


class TestMcNemarTest:
    def test_classic_example(self) -> None:
        """n_01=10, n_10=3 -> p ~0.052 (exact binomial, not significant at 0.05)."""
        pairs: list[BinaryPair] = (
            [(False, True)] * 10  # A fails, B succeeds
            + [(True, False)] * 3  # A succeeds, B fails
            + [(True, True)] * 5  # concordant
            + [(False, False)] * 2  # concordant
        )
        result = mcnemar_test(pairs)
        assert result["n_01"] == 10
        assert result["n_10"] == 3
        assert result["n_discordant"] == 13
        assert result["n_concordant"] == 7
        assert result["method"] == "exact_binomial"
        # Exact binomial p-value for 10 out of 13 with p=0.5
        assert 0.04 < result["p_value"] < 0.10
        # At alpha=0.05, this is borderline
        assert isinstance(result["significant"], bool)

    def test_large_sample_chi_squared(self) -> None:
        """>=25 discordant pairs uses chi-squared with continuity correction."""
        pairs: list[BinaryPair] = [(False, True)] * 20 + [(True, False)] * 5
        result = mcnemar_test(pairs)
        assert result["n_discordant"] == 25
        assert result["method"] == "chi_squared_cc"
        # With 20 vs 5, should be significant
        assert result["p_value"] < 0.01
        assert result["significant"] is True

    def test_all_concordant(self) -> None:
        """All concordant pairs -> p=1.0, no discordant."""
        pairs: list[BinaryPair] = [(True, True)] * 10 + [(False, False)] * 5
        result = mcnemar_test(pairs)
        assert result["n_discordant"] == 0
        assert result["p_value"] == 1.0
        assert result["method"] == "no_discordant_pairs"
        assert result["significant"] is False

    def test_single_discordant_pair(self) -> None:
        """Single discordant pair -> exact binomial on 1 trial."""
        pairs: list[BinaryPair] = [(False, True)]
        result = mcnemar_test(pairs)
        assert result["n_01"] == 1
        assert result["n_10"] == 0
        assert result["n_discordant"] == 1
        assert result["p_value"] == 1.0  # 1 out of 1, two-sided

    def test_empty_pairs(self) -> None:
        """Empty input -> no discordant pairs."""
        result = mcnemar_test([])
        assert result["n_discordant"] == 0
        assert result["p_value"] == 1.0

    def test_symmetric_discordant(self) -> None:
        """Equal discordant both ways -> not significant."""
        pairs: list[BinaryPair] = [(False, True)] * 10 + [(True, False)] * 10
        result = mcnemar_test(pairs)
        assert result["n_01"] == 10
        assert result["n_10"] == 10
        assert result["p_value"] > 0.5


# ---------------------------------------------------------------------------
# Cohen's h
# ---------------------------------------------------------------------------


class TestCohensH:
    def test_standard_example(self) -> None:
        """(0.5, 0.7) -> h ~0.407 (small-medium)."""
        result = cohens_h(0.5, 0.7)
        assert 0.38 < result["h"] < 0.44
        assert result["magnitude"] == "small"

    def test_zero_to_one(self) -> None:
        """(0.0, 1.0) -> h = pi (large)."""
        result = cohens_h(0.0, 1.0)
        assert abs(result["h"] - math.pi) < 0.01
        assert result["magnitude"] == "large"

    def test_equal_proportions(self) -> None:
        """Same proportions -> h = 0 (negligible)."""
        result = cohens_h(0.5, 0.5)
        assert abs(result["h"]) < 0.001
        assert result["magnitude"] == "negligible"

    def test_negative_direction(self) -> None:
        """p1 > p2 -> negative h."""
        result = cohens_h(0.8, 0.3)
        assert result["h"] < -0.5

    def test_clamped_inputs(self) -> None:
        """Values outside [0,1] are clamped."""
        result = cohens_h(-0.1, 1.5)
        # Should be equivalent to (0.0, 1.0)
        assert abs(result["h"] - math.pi) < 0.01


# ---------------------------------------------------------------------------
# Bootstrap CI
# ---------------------------------------------------------------------------


class TestBootstrapCI:
    def test_deterministic_with_seed(self) -> None:
        """Same seed produces same CI."""
        pairs: list[BinaryPair] = [(False, True)] * 15 + [(True, True)] * 5
        r1 = bootstrap_ci(pairs, seed=42, n_bootstrap=1000)
        r2 = bootstrap_ci(pairs, seed=42, n_bootstrap=1000)
        assert r1["ci_lower"] == r2["ci_lower"]
        assert r1["ci_upper"] == r2["ci_upper"]

    def test_ci_contains_observed(self) -> None:
        """CI should generally contain the observed uplift."""
        pairs: list[BinaryPair] = (
            [(False, True)] * 10 + [(True, False)] * 5 + [(True, True)] * 5 + [(False, False)] * 10
        )
        result = bootstrap_ci(pairs, seed=123, n_bootstrap=5000)
        # Observed uplift: (10+5)/30 B-rate - (5+5)/30 A-rate
        assert result["ci_lower"] <= result["uplift"] <= result["ci_upper"]

    def test_empty_pairs(self) -> None:
        """Empty input -> zero uplift, zero CI."""
        result = bootstrap_ci([])
        assert result["uplift"] == 0.0
        assert result["ci_lower"] == 0.0
        assert result["ci_upper"] == 0.0
        assert result["n_bootstrap"] == 0

    def test_positive_uplift(self) -> None:
        """Strong B advantage should give positive CI."""
        pairs: list[BinaryPair] = [(False, True)] * 20 + [(False, False)] * 5
        result = bootstrap_ci(pairs, seed=42, n_bootstrap=5000)
        assert result["uplift"] > 0
        assert result["ci_lower"] > 0  # Entirely positive


# ---------------------------------------------------------------------------
# Holm-Bonferroni
# ---------------------------------------------------------------------------


class TestHolmBonferroni:
    def test_basic_correction(self) -> None:
        """Three p-values with different significance levels."""
        p_values = {"a": 0.01, "b": 0.04, "c": 0.06}
        result = holm_bonferroni(p_values, alpha=0.05)

        # a: p=0.01, compare with 0.05/3 = 0.0167 -> rejected
        assert result["a"]["rejected"] is True
        assert result["a"]["rank"] == 1

        # b: p=0.04, compare with 0.05/2 = 0.025 -> NOT rejected
        assert result["b"]["rejected"] is False
        assert result["b"]["rank"] == 2

        # c: p=0.06, not rejected since b wasn't
        assert result["c"]["rejected"] is False
        assert result["c"]["rank"] == 3

    def test_all_significant(self) -> None:
        """Very small p-values -> all rejected."""
        p_values = {"x": 0.001, "y": 0.002, "z": 0.003}
        result = holm_bonferroni(p_values, alpha=0.05)
        assert all(v["rejected"] for v in result.values())

    def test_none_significant(self) -> None:
        """Large p-values -> none rejected."""
        p_values = {"a": 0.5, "b": 0.6}
        result = holm_bonferroni(p_values, alpha=0.05)
        assert not any(v["rejected"] for v in result.values())

    def test_empty_input(self) -> None:
        """Empty dict -> empty result."""
        assert holm_bonferroni({}) == {}

    def test_single_p_value(self) -> None:
        """Single test -> just compare directly against alpha."""
        result = holm_bonferroni({"only": 0.03}, alpha=0.05)
        assert result["only"]["rejected"] is True

    def test_adjusted_p_values_monotonic(self) -> None:
        """Adjusted p-values should be non-decreasing when sorted by original p."""
        p_values = {"a": 0.01, "b": 0.02, "c": 0.03, "d": 0.04}
        result = holm_bonferroni(p_values)
        sorted_by_rank = sorted(result.values(), key=lambda v: v["rank"])
        adjusted = [v["adjusted_p"] for v in sorted_by_rank]
        for i in range(len(adjusted) - 1):
            assert adjusted[i] <= adjusted[i + 1]


# ---------------------------------------------------------------------------
# Power report
# ---------------------------------------------------------------------------


class TestPowerReport:
    def test_sufficient_power(self) -> None:
        """Large n_discordant with strong uplift -> sufficient."""
        result = power_report(200, 80, 0.4)
        assert result["sufficient"] is True
        assert result["approximate_power"] >= 0.8

    def test_zero_discordant(self) -> None:
        """Zero discordant -> zero power."""
        result = power_report(50, 0, 0.2)
        assert result["approximate_power"] == 0.0
        assert result["sufficient"] is False

    def test_zero_uplift(self) -> None:
        """Zero uplift -> zero power."""
        result = power_report(50, 20, 0.0)
        assert result["approximate_power"] == 0.0
        assert result["sufficient"] is False


# ---------------------------------------------------------------------------
# Majority vote / any-success aggregation
# ---------------------------------------------------------------------------


class TestAggregation:
    def test_majority_vote_basic(self) -> None:
        """Majority vote: 2/3 runs succeed -> True."""
        runs: dict[str, list[BinaryPair]] = {
            "case1": [(True, False), (True, True), (False, True)],
            "case2": [(False, False), (False, False), (False, True)],
        }
        result = majority_vote_aggregation(runs)
        assert len(result) == 2
        # case1: A=2/3 -> True, B=2/3 -> True
        assert result[0] == (True, True)
        # case2: A=0/3 -> False, B=1/3 -> False
        assert result[1] == (False, False)

    def test_any_success_basic(self) -> None:
        """Any-success: any True -> True."""
        runs: dict[str, list[BinaryPair]] = {
            "case1": [(False, False), (True, False)],
            "case2": [(False, False), (False, False)],
        }
        result = any_success_aggregation(runs)
        assert len(result) == 2
        assert result[0] == (True, False)
        assert result[1] == (False, False)

    def test_empty_runs(self) -> None:
        """Empty case runs are skipped."""
        runs: dict[str, list[BinaryPair]] = {"case1": [], "case2": [(True, True)]}
        result = majority_vote_aggregation(runs)
        assert len(result) == 1


# ---------------------------------------------------------------------------
# Full analysis pipeline
# ---------------------------------------------------------------------------


class TestAnalyzePairedResults:
    def test_full_pipeline(self) -> None:
        """End-to-end: analyze_paired_results with matched rows."""
        rows = [
            {
                "run_id": "r1",
                "case_id": "c1",
                "condition": "traceback_only",
                "category": "hidden_state",
                "success": False,
            },
            {
                "run_id": "r1",
                "case_id": "c1",
                "condition": "with_snapshot",
                "category": "hidden_state",
                "success": True,
            },
            {
                "run_id": "r1",
                "case_id": "c2",
                "condition": "traceback_only",
                "category": "parse_value",
                "success": True,
            },
            {
                "run_id": "r1",
                "case_id": "c2",
                "condition": "with_snapshot",
                "category": "parse_value",
                "success": True,
            },
        ]
        result = analyze_paired_results(rows, seed=42)
        assert result["n_pairs"] == 2
        assert result["a_success_rate"] == 0.5
        assert result["b_success_rate"] == 1.0
        assert "overall_mcnemar" in result
        assert "overall_cohens_h" in result
        assert "per_category" in result
        assert "hidden_state" in result["per_category"]

    def test_full_pipeline_with_custom_pair_conditions(self) -> None:
        rows = [
            {
                "run_id": "r1",
                "case_id": "c1",
                "condition": "traceback_only",
                "category": "hidden_state",
                "success": False,
            },
            {
                "run_id": "r1",
                "case_id": "c1",
                "condition": "with_snapshot_structured",
                "category": "hidden_state",
                "success": True,
            },
            {
                "run_id": "r1",
                "case_id": "c2",
                "condition": "traceback_only",
                "category": "parse_value",
                "success": True,
            },
            {
                "run_id": "r1",
                "case_id": "c2",
                "condition": "with_snapshot_structured",
                "category": "parse_value",
                "success": True,
            },
        ]
        result = analyze_paired_results(
            rows,
            seed=42,
            condition_a="traceback_only",
            condition_b="with_snapshot_structured",
        )
        assert result["n_pairs"] == 2
        assert result["condition_a"] == "traceback_only"
        assert result["condition_b"] == "with_snapshot_structured"
        assert result["a_success_rate"] == 0.5
        assert result["b_success_rate"] == 1.0

    def test_no_matched_pairs(self) -> None:
        """Unmatched conditions -> error."""
        rows = [
            {
                "run_id": "r1",
                "case_id": "c1",
                "condition": "traceback_only",
                "category": "hidden_state",
                "success": False,
            },
        ]
        result = analyze_paired_results(rows)
        assert "error" in result

    def test_empty_rows(self) -> None:
        """Empty input -> error."""
        result = analyze_paired_results([])
        assert "error" in result


class TestAnalyzeConditionMatrix:
    def test_matrix_pairs_all_conditions(self) -> None:
        rows = [
            {
                "run_id": "r1",
                "case_id": "c1",
                "condition": "traceback_only",
                "category": "hidden_state",
                "success": False,
            },
            {
                "run_id": "r1",
                "case_id": "c1",
                "condition": "with_snapshot",
                "category": "hidden_state",
                "success": True,
            },
            {
                "run_id": "r1",
                "case_id": "c1",
                "condition": "with_snapshot_structured",
                "category": "hidden_state",
                "success": True,
            },
            {
                "run_id": "r1",
                "case_id": "c2",
                "condition": "traceback_only",
                "category": "parse_value",
                "success": False,
            },
            {
                "run_id": "r1",
                "case_id": "c2",
                "condition": "with_snapshot",
                "category": "parse_value",
                "success": False,
            },
            {
                "run_id": "r1",
                "case_id": "c2",
                "condition": "with_snapshot_structured",
                "category": "parse_value",
                "success": True,
            },
        ]
        matrix = analyze_condition_matrix(rows, seed=7, n_bootstrap=100)
        assert len(matrix) == 3
        pairs = {(item.get("condition_a"), item.get("condition_b")) for item in matrix}
        assert ("traceback_only", "with_snapshot") in pairs
        assert ("traceback_only", "with_snapshot_structured") in pairs
        assert ("with_snapshot", "with_snapshot_structured") in pairs
