"""Tests for PC algorithm causal discovery."""

from __future__ import annotations

import random

from orxaq.intelligence.pc_algorithm import PCAlgorithm


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


def _make_fork_data(n: int = 500, seed: int = 42) -> list[dict]:
    """X <- Z -> Y fork."""
    rng = random.Random(seed)
    data = []
    for _ in range(n):
        z = rng.gauss(0, 1)
        x = 2.0 * z + rng.gauss(0, 0.3)
        y = 1.5 * z + rng.gauss(0, 0.3)
        data.append({"X": x, "Z": z, "Y": y})
    return data


def _make_collider_data(n: int = 500, seed: int = 42) -> list[dict]:
    """X -> Z <- Y collider."""
    rng = random.Random(seed)
    data = []
    for _ in range(n):
        x = rng.gauss(0, 1)
        y = rng.gauss(0, 1)
        z = 2.0 * x + 1.5 * y + rng.gauss(0, 0.3)
        data.append({"X": x, "Z": z, "Y": y})
    return data


class TestPCChain:
    def test_recovers_chain_skeleton(self):
        """PC should find edges X-Z and Z-Y in a chain."""
        data = _make_chain_data(n=800)
        pc = PCAlgorithm(alpha=0.05)
        result = pc.fit(data, ["X", "Z", "Y"])

        # Skeleton should contain X-Z and Z-Y
        skeleton_set = set()
        for a, b in result.skeleton:
            skeleton_set.add((a, b))
            skeleton_set.add((b, a))

        assert ("X", "Z") in skeleton_set or ("Z", "X") in skeleton_set
        assert ("Z", "Y") in skeleton_set or ("Y", "Z") in skeleton_set

    def test_no_direct_xy_edge(self):
        """X and Y should NOT be directly connected (separated by Z)."""
        data = _make_chain_data(n=1000)
        pc = PCAlgorithm(alpha=0.05)
        result = pc.fit(data, ["X", "Z", "Y"])

        # X-Y should not be in skeleton
        for a, b in result.skeleton:
            pair = tuple(sorted([a, b]))
            assert pair != ("X", "Y"), "X-Y edge should not exist in chain"


class TestPCFork:
    def test_recovers_fork_skeleton(self):
        """PC should find edges X-Z and Z-Y in a fork."""
        data = _make_fork_data(n=800)
        pc = PCAlgorithm(alpha=0.05)
        result = pc.fit(data, ["X", "Z", "Y"])

        skeleton_set = set()
        for a, b in result.skeleton:
            skeleton_set.add((a, b))
            skeleton_set.add((b, a))

        assert ("X", "Z") in skeleton_set or ("Z", "X") in skeleton_set
        assert ("Z", "Y") in skeleton_set or ("Y", "Z") in skeleton_set


class TestPCCollider:
    def test_recovers_collider_skeleton(self):
        """PC should find edges X-Z and Y-Z in a collider."""
        data = _make_collider_data(n=800)
        pc = PCAlgorithm(alpha=0.05)
        result = pc.fit(data, ["X", "Z", "Y"])

        skeleton_set = set()
        for a, b in result.skeleton:
            skeleton_set.add((a, b))
            skeleton_set.add((b, a))

        assert ("X", "Z") in skeleton_set or ("Z", "X") in skeleton_set
        assert ("Y", "Z") in skeleton_set or ("Z", "Y") in skeleton_set

    def test_orients_collider(self):
        """PC should orient as X->Z<-Y."""
        data = _make_collider_data(n=1000)
        pc = PCAlgorithm(alpha=0.05)
        result = pc.fit(data, ["X", "Z", "Y"])

        # Should have orientations into Z
        assert ("X", "Z") in result.oriented or ("Y", "Z") in result.oriented


class TestPCConstraints:
    def test_respects_forbidden_edges(self):
        """Forbidden edges should not appear in skeleton."""
        data = _make_chain_data(n=500)
        pc = PCAlgorithm(alpha=0.05)
        result = pc.fit(data, ["X", "Z", "Y"], forbidden_edges={("X", "Z"), ("Z", "X")})

        for a, b in result.skeleton:
            assert not (a == "X" and b == "Z")
            assert not (a == "Z" and b == "X")

    def test_respects_required_edges(self):
        """Required edges should always be in skeleton."""
        data = _make_chain_data(n=500)
        pc = PCAlgorithm(alpha=0.05)
        result = pc.fit(data, ["X", "Z", "Y"], required_edges={("X", "Y")})

        skeleton_pairs = set()
        for a, b in result.skeleton:
            skeleton_pairs.add(tuple(sorted([a, b])))
        assert ("X", "Y") in skeleton_pairs

    def test_records_separating_sets(self):
        """Separating sets should be recorded for removed edges."""
        data = _make_chain_data(n=1000)
        pc = PCAlgorithm(alpha=0.05)
        result = pc.fit(data, ["X", "Z", "Y"])

        # If X-Y was removed, its separating set should contain Z
        key = ("X", "Y")
        if key in result.separating_sets:
            assert "Z" in result.separating_sets[key]

    def test_ci_tests_counted(self):
        """ci_tests_run should be positive."""
        data = _make_chain_data(n=500)
        pc = PCAlgorithm(alpha=0.05)
        result = pc.fit(data, ["X", "Z", "Y"])
        assert result.ci_tests_run > 0


# ── Coverage: non-numeric data handling (lines 117-118) ──────


class TestPCNonNumericData:
    """Cover ValueError/TypeError handling in _skeleton."""

    def test_non_numeric_values_handled(self):
        """Non-numeric values should be treated as 0.0."""
        data = [
            {"X": 1.0, "Z": "bad", "Y": 2.0},
            {"X": 2.0, "Z": None, "Y": 4.0},
            {"X": 3.0, "Z": 6.0, "Y": 6.0},
        ] * 100
        pc = PCAlgorithm(alpha=0.05)
        result = pc.fit(data, ["X", "Z", "Y"])
        assert result.ci_tests_run > 0

    def test_missing_keys_handled(self):
        """Missing keys should be treated as 0.0."""
        data = [
            {"X": 1.0, "Y": 2.0},  # Z missing
            {"X": 2.0, "Z": 4.0, "Y": 4.0},
        ] * 100
        pc = PCAlgorithm(alpha=0.05)
        result = pc.fit(data, ["X", "Z", "Y"])
        assert result.ci_tests_run > 0


# ── Coverage: required edge skip in skeleton (line 136) ──────


class TestPCRequiredEdgesInSkeleton:
    """Cover required edge skip in skeleton phase."""

    def test_required_edge_preserved_in_skeleton(self):
        """Required edges skip CI testing and remain in skeleton."""
        data = _make_chain_data(n=1000)
        pc = PCAlgorithm(alpha=0.05)
        result = pc.fit(
            data,
            ["X", "Z", "Y"],
            required_edges={("X", "Y")},
        )
        skeleton_pairs = {tuple(sorted([a, b])) for a, b in result.skeleton}
        assert ("X", "Y") in skeleton_pairs


# ── Coverage: Meek's Rule 2 and Rule 3 (lines 226-244) ──────


class TestPCOrientationRules:
    """Cover Meek's Rule 2 and Rule 3."""

    def test_rule2_propagates_orientation(self):
        """With X->Z<-W collider and Z-Y, Rule 2 orients Z->Y."""
        rng = random.Random(42)
        data = []
        for _ in range(1000):
            x = rng.gauss(0, 1)
            w = rng.gauss(0, 1)
            z = 2.0 * x + 2.0 * w + rng.gauss(0, 0.1)
            y = 2.0 * z + rng.gauss(0, 0.1)
            data.append({"X": x, "W": w, "Z": z, "Y": y})
        pc = PCAlgorithm(alpha=0.01)
        result = pc.fit(data, ["X", "W", "Z", "Y"])
        # With collider X->Z<-W, Rule 2 should orient Z->Y
        if ("X", "Z") in result.oriented or ("W", "Z") in result.oriented:
            assert ("Z", "Y") in result.oriented or len(result.oriented) > 2

    def test_rule3_diamond_propagation(self):
        """Rule 3: A->C, C->D, A-D -> A->D (diamond structure).

        Graph: A->C<-B (collider), C->D, A->D.
        After collider: A->C, B->C oriented.
        Rule 2: B->C, C-D, B not adj D -> C->D.
        Rule 3: A->C, C->D, A-D -> A->D.
        """
        rng = random.Random(42)
        data = []
        for _ in range(1500):
            a = rng.gauss(0, 1)
            b = rng.gauss(0, 1)
            c = 3.0 * a + 3.0 * b + rng.gauss(0, 0.1)
            d = 2.0 * c + 2.0 * a + rng.gauss(0, 0.1)
            data.append({"A": a, "B": b, "C": c, "D": d})
        pc = PCAlgorithm(alpha=0.01, max_cond_size=3)
        result = pc.fit(data, ["A", "B", "C", "D"])
        # Should orient at least the collider edges
        assert len(result.oriented) >= 2
