"""Tests for compute_party_bootstrap_cis in analysis.political_axis."""

import numpy as np

from analysis.political_axis import compute_party_bootstrap_cis


# ── Helpers ──────────────────────────────────────────────────────────────────


def _make_party_vectors(n_mps: int, dim: int = 50, seed: int = 0) -> list:
    """Generate a list of random MP vectors for a single party."""
    rng = np.random.default_rng(seed)
    return [rng.standard_normal(dim) for _ in range(n_mps)]


# ── Tests ────────────────────────────────────────────────────────────────────


class TestBootstrapDeterministic:
    def test_same_seed_gives_identical_output(self):
        """Same inputs + same seed -> identical outputs."""
        vecs = _make_party_vectors(10, dim=5, seed=99)
        party_vectors = {"PartyA": vecs}

        result1 = compute_party_bootstrap_cis(party_vectors, n_boot=200, seed=42)
        result2 = compute_party_bootstrap_cis(party_vectors, n_boot=200, seed=42)

        np.testing.assert_array_equal(
            result1["PartyA"]["centroid"], result2["PartyA"]["centroid"]
        )
        np.testing.assert_array_equal(
            result1["PartyA"]["ci_lower"], result2["PartyA"]["ci_lower"]
        )
        np.testing.assert_array_equal(
            result1["PartyA"]["ci_upper"], result2["PartyA"]["ci_upper"]
        )
        np.testing.assert_array_equal(
            result1["PartyA"]["std"], result2["PartyA"]["std"]
        )
        assert result1["PartyA"]["n_mps"] == result2["PartyA"]["n_mps"]


class TestBootstrapSingleMP:
    def test_single_mp_collapses_ci(self):
        """Party with 1 MP -> ci_lower == ci_upper == centroid, std == 0."""
        vec = np.array([1.0, 2.0, 3.0])
        party_vectors = {"Solo": [vec]}

        result = compute_party_bootstrap_cis(party_vectors, n_boot=500)
        entry = result["Solo"]

        np.testing.assert_array_equal(entry["centroid"], vec)
        np.testing.assert_array_equal(entry["ci_lower"], vec)
        np.testing.assert_array_equal(entry["ci_upper"], vec)
        np.testing.assert_array_equal(entry["std"], np.zeros_like(vec))
        assert entry["n_mps"] == 1


class TestBootstrapCIWidthScalesWithN:
    def test_larger_party_has_narrower_ci(self):
        """Party with 3 MPs should have wider CIs than party with 30 MPs
        when sampled from the same distribution."""
        rng = np.random.default_rng(123)
        dim = 10
        # Same underlying distribution, different sample sizes
        small_vecs = [rng.standard_normal(dim) for _ in range(3)]
        large_vecs = [rng.standard_normal(dim) for _ in range(30)]

        party_vectors = {"Small": small_vecs, "Large": large_vecs}
        result = compute_party_bootstrap_cis(party_vectors, n_boot=2000, seed=42)

        small_width = result["Small"]["ci_upper"] - result["Small"]["ci_lower"]
        large_width = result["Large"]["ci_upper"] - result["Large"]["ci_lower"]

        # On average, the small party's CI should be wider
        assert np.mean(small_width) > np.mean(large_width)


class TestBootstrapEmptyParty:
    def test_empty_list_excluded(self):
        """Party with empty list -> excluded from output."""
        party_vectors = {
            "HasMPs": _make_party_vectors(5, dim=4),
            "Empty": [],
        }

        result = compute_party_bootstrap_cis(party_vectors, n_boot=100)

        assert "HasMPs" in result
        assert "Empty" not in result


class TestBootstrapCIContainsCentroid:
    def test_centroid_within_ci_bounds(self):
        """ci_lower <= centroid <= ci_upper for each dimension."""
        party_vectors = {"A": _make_party_vectors(15, dim=8, seed=7)}
        result = compute_party_bootstrap_cis(party_vectors, n_boot=1000, seed=42)

        entry = result["A"]
        assert np.all(entry["ci_lower"] <= entry["centroid"])
        assert np.all(entry["centroid"] <= entry["ci_upper"])


class TestBootstrapCustomCILevel:
    def test_wider_ci_at_higher_level(self):
        """ci=99 produces wider intervals than ci=90."""
        party_vectors = {"X": _make_party_vectors(20, dim=6, seed=55)}

        result_90 = compute_party_bootstrap_cis(
            party_vectors, n_boot=2000, ci=90.0, seed=42
        )
        result_99 = compute_party_bootstrap_cis(
            party_vectors, n_boot=2000, ci=99.0, seed=42
        )

        width_90 = result_90["X"]["ci_upper"] - result_90["X"]["ci_lower"]
        width_99 = result_99["X"]["ci_upper"] - result_99["X"]["ci_lower"]

        # 99% CI should be wider than 90% CI on every dimension
        assert np.all(width_99 >= width_90)