"""Tests for WorldDAG — causal graph, d-separation, topological sort."""

import pytest

from orxaq.kernel.types import CausalEdge, CausalRole, CycleError, Entity
from orxaq.kernel.world_dag import WorldDAG


def _entity(eid: str, role: CausalRole = CausalRole.EXOGENOUS) -> Entity:
    """Helper to create a minimal entity."""
    return Entity(entity_id=eid, name=eid, causal_role=role)


def _edge(src: str, tgt: str) -> CausalEdge:
    """Helper to create a minimal edge."""
    return CausalEdge(source=src, target=tgt)


class TestBasicOperations:
    def test_add_entity_and_edge(self):
        dag = WorldDAG()
        dag.add_entity(_entity("A"))
        dag.add_entity(_entity("B"))
        dag.add_edge(_edge("A", "B"))
        assert dag.node_count() == 2
        assert dag.edge_count() == 1

    def test_cycle_detection(self):
        dag = WorldDAG()
        dag.add_entity(_entity("A"))
        dag.add_entity(_entity("B"))
        dag.add_edge(_edge("A", "B"))
        with pytest.raises(CycleError):
            dag.add_edge(_edge("B", "A"))

    def test_self_loop_rejected(self):
        dag = WorldDAG()
        dag.add_entity(_entity("A"))
        with pytest.raises(CycleError, match="Self-loop"):
            dag.add_edge(_edge("A", "A"))

    def test_parents_and_children(self):
        dag = WorldDAG()
        dag.add_entity(_entity("A"))
        dag.add_entity(_entity("B"))
        dag.add_entity(_entity("C"))
        dag.add_edge(_edge("A", "B"))
        dag.add_edge(_edge("A", "C"))
        assert dag.children("A") == {"B", "C"}
        assert dag.parents("B") == {"A"}
        assert dag.parents("A") == set()

    def test_ancestors_and_descendants(self):
        dag = WorldDAG()
        for eid in ["A", "B", "C", "D"]:
            dag.add_entity(_entity(eid))
        dag.add_edge(_edge("A", "B"))
        dag.add_edge(_edge("B", "C"))
        dag.add_edge(_edge("C", "D"))
        assert dag.ancestors("D") == {"A", "B", "C"}
        assert dag.descendants("A") == {"B", "C", "D"}

    def test_topological_sort_respects_ordering(self):
        dag = WorldDAG()
        for eid in ["A", "B", "C"]:
            dag.add_entity(_entity(eid))
        dag.add_edge(_edge("A", "B"))
        dag.add_edge(_edge("B", "C"))
        order = dag.topological_sort()
        assert order.index("A") < order.index("B") < order.index("C")

    def test_all_directed_paths(self):
        dag = WorldDAG()
        for eid in ["A", "B", "C", "D"]:
            dag.add_entity(_entity(eid))
        dag.add_edge(_edge("A", "B"))
        dag.add_edge(_edge("A", "C"))
        dag.add_edge(_edge("B", "D"))
        dag.add_edge(_edge("C", "D"))
        paths = dag.all_directed_paths("A", "D")
        assert len(paths) == 2
        path_strs = [" ".join(p) for p in paths]
        assert "A B D" in path_strs or "A C D" in path_strs

    def test_remove_entity_cascades_edges(self):
        dag = WorldDAG()
        dag.add_entity(_entity("A"))
        dag.add_entity(_entity("B"))
        dag.add_entity(_entity("C"))
        dag.add_edge(_edge("A", "B"))
        dag.add_edge(_edge("B", "C"))
        dag.remove_entity("B")
        assert dag.node_count() == 2
        assert dag.edge_count() == 0
        assert dag.children("A") == set()

    def test_summary(self):
        dag = WorldDAG()
        dag.add_entity(_entity("A"))
        dag.add_entity(_entity("B"))
        dag.add_edge(_edge("A", "B"))
        s = dag.summary()
        assert s == {"nodes": 2, "edges": 1}


class TestDSeparation:
    """Test the three fundamental causal structures."""

    def _build_chain(self) -> WorldDAG:
        """Chain: X → Z → Y"""
        dag = WorldDAG()
        for eid in ["X", "Z", "Y"]:
            dag.add_entity(_entity(eid))
        dag.add_edge(_edge("X", "Z"))
        dag.add_edge(_edge("Z", "Y"))
        return dag

    def _build_fork(self) -> WorldDAG:
        """Fork: X ← Z → Y"""
        dag = WorldDAG()
        for eid in ["X", "Z", "Y"]:
            dag.add_entity(_entity(eid))
        dag.add_edge(_edge("Z", "X"))
        dag.add_edge(_edge("Z", "Y"))
        return dag

    def _build_collider(self) -> WorldDAG:
        """Collider: X → Z ← Y"""
        dag = WorldDAG()
        for eid in ["X", "Z", "Y"]:
            dag.add_entity(_entity(eid))
        dag.add_edge(_edge("X", "Z"))
        dag.add_edge(_edge("Y", "Z"))
        return dag

    def test_chain_blocked_when_conditioned(self):
        """X → Z → Y: X ⊥ Y | Z"""
        dag = self._build_chain()
        assert dag.d_separated("X", "Y", frozenset({"Z"})) is True

    def test_chain_open_when_unconditioned(self):
        """X → Z → Y: X ⊥̸ Y | ∅"""
        dag = self._build_chain()
        assert dag.d_separated("X", "Y", frozenset()) is False

    def test_fork_blocked_when_conditioned(self):
        """X ← Z → Y: X ⊥ Y | Z"""
        dag = self._build_fork()
        assert dag.d_separated("X", "Y", frozenset({"Z"})) is True

    def test_fork_open_when_unconditioned(self):
        """X ← Z → Y: X ⊥̸ Y | ∅"""
        dag = self._build_fork()
        assert dag.d_separated("X", "Y", frozenset()) is False

    def test_collider_blocked_when_unconditioned(self):
        """X → Z ← Y: X ⊥ Y | ∅"""
        dag = self._build_collider()
        assert dag.d_separated("X", "Y", frozenset()) is True

    def test_collider_open_when_conditioned(self):
        """X → Z ← Y: X ⊥̸ Y | Z"""
        dag = self._build_collider()
        assert dag.d_separated("X", "Y", frozenset({"Z"})) is False


class TestBackdoorAdjustment:
    def test_adjustment_with_confounder(self):
        """Z → X, Z → Y, X → Y: Z is a confounder, must be adjusted."""
        dag = WorldDAG()
        for eid in ["Z", "X", "Y"]:
            dag.add_entity(_entity(eid))
        dag.add_edge(_edge("Z", "X"))
        dag.add_edge(_edge("Z", "Y"))
        dag.add_edge(_edge("X", "Y"))
        adj_set = dag.backdoor_adjustment_set("X", "Y")
        assert "Z" in adj_set
