"""World DAG — the causal graph at the heart of Orxaq.

Implements d-separation via the Bayes-Ball algorithm, topological sort
via Kahn's algorithm, cycle detection, and backdoor adjustment.

All stdlib only (ADR-007).
"""

from __future__ import annotations

from collections import deque

from orxaq.kernel.types import CausalEdge, CycleError, Entity


class WorldDAG:
    """Directed acyclic graph for causal reasoning.

    Dict-based adjacency for O(1) lookups. Entities are nodes,
    CausalEdges are directed edges with causal metadata.
    """

    def __init__(self) -> None:
        self._entities: dict[str, Entity] = {}
        self._edges: dict[str, CausalEdge] = {}
        self._children_map: dict[str, set[str]] = {}
        self._parents_map: dict[str, set[str]] = {}

    def add_entity(self, entity: Entity) -> None:
        """Add an entity (node) to the DAG."""
        self._entities[entity.entity_id] = entity
        self._children_map.setdefault(entity.entity_id, set())
        self._parents_map.setdefault(entity.entity_id, set())

    def get_entity(self, entity_id: str) -> Entity | None:
        """Retrieve an entity by ID, or None if not found."""
        return self._entities.get(entity_id)

    def remove_entity(self, entity_id: str) -> None:
        """Remove an entity and all edges connected to it."""
        if entity_id not in self._entities:
            return

        edges_to_remove = []
        for edge_key, edge in self._edges.items():
            if edge.source == entity_id or edge.target == entity_id:
                edges_to_remove.append(edge_key)

        for edge_key in edges_to_remove:
            edge = self._edges[edge_key]
            self._children_map.get(edge.source, set()).discard(edge.target)
            self._parents_map.get(edge.target, set()).discard(edge.source)
            del self._edges[edge_key]

        del self._entities[entity_id]
        self._children_map.pop(entity_id, None)
        self._parents_map.pop(entity_id, None)

    def add_edge(self, edge: CausalEdge) -> None:
        """Add a causal edge. Raises CycleError if it would create a cycle."""
        if edge.source == edge.target:
            raise CycleError(f"Self-loop: {edge.source} \u2192 {edge.target}")

        if edge.source not in self._entities:
            raise ValueError(f"Source entity '{edge.source}' not in DAG")
        if edge.target not in self._entities:
            raise ValueError(f"Target entity '{edge.target}' not in DAG")

        if self._would_create_cycle(edge.source, edge.target):
            raise CycleError(f"Cycle detected: {edge.source} \u2192 {edge.target}")

        self._edges[edge.edge_id] = edge
        self._children_map[edge.source].add(edge.target)
        self._parents_map[edge.target].add(edge.source)

    def _would_create_cycle(self, source: str, target: str) -> bool:
        """Check if adding source→target would create a cycle (target can reach source)."""
        visited: set[str] = set()
        queue = deque([target])
        while queue:
            node = queue.popleft()
            if node == source:
                return True
            if node in visited:
                continue
            visited.add(node)
            queue.extend(self._children_map.get(node, set()))
        return False

    def parents(self, entity_id: str) -> set[str]:
        """Direct parents of an entity."""
        return set(self._parents_map.get(entity_id, set()))

    def children(self, entity_id: str) -> set[str]:
        """Direct children of an entity."""
        return set(self._children_map.get(entity_id, set()))

    def ancestors(self, entity_id: str) -> set[str]:
        """All ancestors (transitive parents) of an entity."""
        result: set[str] = set()
        queue = deque(self._parents_map.get(entity_id, set()))
        while queue:
            node = queue.popleft()
            if node in result:
                continue
            result.add(node)
            queue.extend(self._parents_map.get(node, set()))
        return result

    def descendants(self, entity_id: str) -> set[str]:
        """All descendants (transitive children) of an entity."""
        result: set[str] = set()
        queue = deque(self._children_map.get(entity_id, set()))
        while queue:
            node = queue.popleft()
            if node in result:
                continue
            result.add(node)
            queue.extend(self._children_map.get(node, set()))
        return result

    def topological_sort(self) -> list[str]:
        """Kahn's algorithm for topological ordering."""
        in_degree: dict[str, int] = {eid: 0 for eid in self._entities}
        for eid in self._entities:
            for child in self._children_map.get(eid, set()):
                in_degree[child] += 1

        queue = deque(eid for eid, deg in in_degree.items() if deg == 0)
        result: list[str] = []

        while queue:
            node = queue.popleft()
            result.append(node)
            for child in sorted(self._children_map.get(node, set())):
                in_degree[child] -= 1
                if in_degree[child] == 0:
                    queue.append(child)

        return result

    def all_directed_paths(self, source: str, target: str) -> list[list[str]]:
        """Find all directed paths from source to target via DFS."""
        paths: list[list[str]] = []
        stack: list[tuple[str, list[str]]] = [(source, [source])]

        while stack:
            node, path = stack.pop()
            if node == target:
                paths.append(path)
                continue
            for child in sorted(self._children_map.get(node, set())):
                if child not in path:
                    stack.append((child, path + [child]))

        return paths

    def d_separated(self, x: str, y: str, z: frozenset[str]) -> bool:
        """Test d-separation using the Bayes-Ball algorithm.

        Returns True if X is d-separated from Y given Z.
        The three fundamental structures:
        - Chain:    X→Z→Y  — X ⊥ Y | Z (blocked)
        - Fork:     X←Z→Y  — X ⊥ Y | Z (blocked)
        - Collider: X→Z←Y  — X ⊥ Y | ∅ (blocked), X ⊥̸ Y | Z (open)
        """
        reachable = self._bayes_ball_reachable(x, z)
        return y not in reachable

    def _bayes_ball_reachable(self, source: str, conditioned: frozenset[str]) -> set[str]:
        """Bayes-Ball: find all nodes reachable from source given conditioning set.

        Uses Shachter (1998) algorithm. Queue entries are (node, from_child)
        where from_child=True means the ball is traveling upward (arrived from child).
        """
        conditioned_set = set(conditioned)

        # Precompute ancestors of ALL conditioned nodes in a single BFS
        # (avoids calling self.ancestors() per conditioned node)
        has_conditioned_descendant: set[str] = set()
        ancestor_queue: deque[str] = deque()
        for c in conditioned_set:
            has_conditioned_descendant.add(c)
            ancestor_queue.extend(self._parents_map.get(c, set()))
        while ancestor_queue:
            node = ancestor_queue.popleft()
            if node in has_conditioned_descendant:
                continue
            has_conditioned_descendant.add(node)
            ancestor_queue.extend(self._parents_map.get(node, set()))

        visited_from_child: set[str] = set()
        visited_from_parent: set[str] = set()
        reachable: set[str] = set()

        # Initialize: send balls from source in BOTH directions
        queue: deque[tuple[str, bool]] = deque()
        for parent in self._parents_map.get(source, set()):
            queue.append((parent, True))  # ball goes up to parent (arrived from child)
        for child in self._children_map.get(source, set()):
            queue.append((child, False))  # ball goes down to child (arrived from parent)

        while queue:
            node, from_child = queue.popleft()

            if from_child and node not in conditioned_set:
                # Ball traveling up, arrives at unconditioned node
                if node not in visited_from_child:
                    visited_from_child.add(node)
                    reachable.add(node)
                    # Continue up to parents
                    for parent in self._parents_map.get(node, set()):
                        queue.append((parent, True))
                    # Bounce down to children
                    for child in self._children_map.get(node, set()):
                        queue.append((child, False))

            elif from_child and node in conditioned_set:
                # Ball traveling up, arrives at conditioned node — blocked
                # (no further propagation)
                pass

            elif not from_child and node not in conditioned_set:
                # Ball traveling down, arrives at unconditioned non-collider
                if node not in visited_from_parent:
                    visited_from_parent.add(node)
                    reachable.add(node)
                    # Continue down to children
                    for child in self._children_map.get(node, set()):
                        queue.append((child, False))

            elif not from_child and node in conditioned_set:
                # Ball traveling down, arrives at conditioned node
                # This opens the collider — bounce up to parents
                if node not in visited_from_parent:
                    visited_from_parent.add(node)
                    for parent in self._parents_map.get(node, set()):
                        queue.append((parent, True))

        reachable.discard(source)
        return reachable

    def backdoor_adjustment_set(self, treatment: str, outcome: str) -> set[str]:
        """Find a sufficient backdoor adjustment set.

        Returns the set of parents of treatment that are also ancestors of outcome
        (confounders that need to be controlled for).
        """
        treatment_parents = self.parents(treatment)
        outcome_ancestors = self.ancestors(outcome)

        adjustment_set: set[str] = set()
        for parent in treatment_parents:
            if parent in outcome_ancestors or parent == outcome:
                adjustment_set.add(parent)
            elif self.descendants(parent) & outcome_ancestors:  # pragma: no cover
                adjustment_set.add(parent)  # pragma: no cover

        return adjustment_set

    @property
    def entities(self) -> dict[str, Entity]:
        """Read-only view of all entities keyed by ID."""
        return dict(self._entities)

    @property
    def edges(self) -> dict[str, CausalEdge]:
        """Read-only view of all edges keyed by edge_id."""
        return dict(self._edges)

    def entity_ids(self) -> list[str]:
        """Return all entity IDs."""
        return list(self._entities.keys())

    def edge_ids(self) -> list[str]:
        """Return all edge IDs."""
        return list(self._edges.keys())

    def node_count(self) -> int:
        """Number of entities in the DAG."""
        return len(self._entities)

    def edge_count(self) -> int:
        """Number of edges in the DAG."""
        return len(self._edges)

    def summary(self) -> dict[str, int]:
        """Summary statistics of the DAG."""
        return {
            "nodes": self.node_count(),
            "edges": self.edge_count(),
        }
