"""Tests for GES algorithm causal discovery."""

from __future__ import annotations

import random

from orxaq.intelligence.ges_algorithm import GESAlgorithm, _ols_rss


def _make_chain_data(n: int = 500, seed: int = 42) -> list[dict]:
    """X → Z → Y chain with strong effects."""
    rng = random.Random(seed)
    data = []
    for _ in range(n):
        x = rng.gauss(0, 1)
        z = 3.0 * x + rng.gauss(0, 0.3)
        y = 3.0 * z + rng.gauss(0, 0.3)
        data.append({"X": x, "Z": z, "Y": y})
    return data


def _make_independent_data(n: int = 500, seed: int = 42) -> list[dict]:
    """Three independent variables."""
    rng = random.Random(seed)
    data = []
    for _ in range(n):
        data.append(
            {
                "X": rng.gauss(0, 1),
                "Y": rng.gauss(0, 1),
                "Z": rng.gauss(0, 1),
            }
        )
    return data


class TestGESChain:
    def test_recovers_chain_edges(self):
        """GES should find edges between causally related variables."""
        data = _make_chain_data(n=500)
        ges = GESAlgorithm()
        result = ges.fit(data, ["X", "Z", "Y"])

        # Should find some edges
        assert len(result.edges) > 0

        # Should find X→Z or Z→X connection
        edge_vars = set()
        for src, tgt in result.edges:
            edge_vars.add(src)
            edge_vars.add(tgt)

        assert "X" in edge_vars
        assert "Z" in edge_vars
        assert "Y" in edge_vars

    def test_bic_score_computed(self):
        """BIC score should be a finite number."""
        data = _make_chain_data(n=200)
        ges = GESAlgorithm()
        result = ges.fit(data, ["X", "Z", "Y"])
        assert result.bic_score != 0.0
        assert not (result.bic_score != result.bic_score)  # not NaN

    def test_additions_counted(self):
        """Forward phase should add edges."""
        data = _make_chain_data(n=200)
        ges = GESAlgorithm()
        result = ges.fit(data, ["X", "Z", "Y"])
        assert result.additions > 0


class TestGESBackward:
    def test_backward_removes_spurious_edges(self):
        """Backward phase should not find too many edges in independent data."""
        data = _make_independent_data(n=500)
        ges = GESAlgorithm()
        result = ges.fit(data, ["X", "Y", "Z"])

        # With independent data, shouldn't have many edges
        assert len(result.edges) <= 2


class TestGESConstraints:
    def test_respects_forbidden_edges(self):
        """Forbidden edges should not appear in result."""
        data = _make_chain_data(n=300)
        ges = GESAlgorithm()
        forbidden = {("X", "Z"), ("Z", "X")}
        result = ges.fit(data, ["X", "Z", "Y"], forbidden_edges=forbidden)

        for src, tgt in result.edges:
            assert (src, tgt) not in forbidden

    def test_respects_required_edges(self):
        """Required edges should appear in result."""
        data = _make_independent_data(n=300)
        ges = GESAlgorithm()
        required = {("X", "Y")}
        result = ges.fit(data, ["X", "Y", "Z"], required_edges=required)

        assert ("X", "Y") in result.edges

    def test_small_data_returns_empty(self):
        """With very small data, should return required edges only."""
        data = [{"X": 1.0, "Y": 2.0, "Z": 3.0}]
        ges = GESAlgorithm()
        result = ges.fit(data, ["X", "Y", "Z"])
        assert len(result.edges) == 0


# ── Direct tests for private scoring function ──────────────


def _make_confounded_data(n: int = 50, seed: int = 5) -> list[dict]:
    """X confounds A and B; C = A + B + noise.  Triggers backward removal."""
    rng = random.Random(seed)
    data = []
    for _ in range(n):
        x = rng.gauss(0, 1)
        a = x + 0.5 * rng.gauss(0, 1)
        b = x + 0.5 * rng.gauss(0, 1)
        c = a + b + rng.gauss(0, 1)
        data.append({"X": x, "A": a, "B": b, "C": c})
    return data


class TestGESBackwardRemoval:
    """Cover lines 114-116, 120-123 (backward phase removal)."""

    def test_backward_actively_removes(self):
        """Confounded data triggers backward removal of spurious edge."""
        data = _make_confounded_data(n=50, seed=5)
        ges = GESAlgorithm()
        result = ges.fit(data, ["X", "A", "B", "C"])
        assert result.additions > 0
        assert result.removals > 0


class TestGESNonNumericExtract:
    """Cover lines 146-147 (ValueError/TypeError in _extract_columns)."""

    def test_string_values_fallback(self):
        """Non-numeric string values should fall back to 0.0."""
        data = [
            {"X": "bad", "Y": 1.0},
            {"X": 2.0, "Y": "bad"},
        ] * 10
        ges = GESAlgorithm()
        result = ges.fit(data, ["X", "Y"])
        assert isinstance(result.bic_score, float)

    def test_object_values_fallback(self):
        """Non-float-coercible values should fall back to 0.0."""
        data = [{"X": [1, 2], "Y": 1.0}] * 10
        ges = GESAlgorithm()
        result = ges.fit(data, ["X", "Y"])
        assert isinstance(result.bic_score, float)

    def test_none_and_missing_values(self):
        """None values and missing keys should fall back to 0.0."""
        data = [{"X": None, "Y": 1.0}, {"Y": 2.0}] * 5
        ges = GESAlgorithm()
        result = ges.fit(data, ["X", "Y"])
        assert isinstance(result.bic_score, float)


class TestGESCycleVisited:
    """Cover line 167 (visited-node skip in cycle detection BFS)."""

    def test_diamond_graph_revisit(self):
        """Diamond graph forces BFS to dequeue an already-visited node."""
        children_map: dict[str, set[str]] = {
            "A": {"B", "C"},
            "B": {"D"},
            "C": {"D"},
            "D": set(),
        }
        result = GESAlgorithm._would_cycle("X", "A", children_map)
        assert result is False


class TestOlsRssEdgeCases:
    """Cover lines 261-263 (_ols_rss k=0) and 282-285 (singular matrix)."""

    def test_no_parents(self):
        """_ols_rss with empty parents computes total SS."""
        col_vals = {"Y": [1.0, 2.0, 3.0, 4.0, 5.0]}
        rss = _ols_rss(col_vals, "Y", [], 5)
        assert abs(rss - 10.0) < 1e-10

    def test_singular_matrix_fallback(self):
        """Linearly dependent parents trigger singular-matrix fallback."""
        n = 20
        col_vals = {
            "T": [float(i * i) for i in range(n)],
            "Y": [float(i) for i in range(n)],
            "Z": [float(2 * i) for i in range(n)],
        }
        rss = _ols_rss(col_vals, "T", ["Y", "Z"], n)
        mean_t = sum(col_vals["T"]) / n
        expected_ss = sum((t - mean_t) ** 2 for t in col_vals["T"])
        assert abs(rss - expected_ss) < 1e-6
