"""PC algorithm — constraint-based causal discovery.

Discovers causal structure by testing conditional independencies.
Phase 1 learns the skeleton, Phase 2 orients edges via Meek's rules.
Stdlib-only (ADR-007).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from itertools import combinations

from orxaq.intelligence.independence import PartialCorrelationTest


@dataclass(frozen=True)
class PCResult:
    """Result of PC algorithm discovery."""

    skeleton: frozenset[tuple[str, str]]
    oriented: frozenset[tuple[str, str]]
    separating_sets: dict[tuple[str, str], frozenset[str]]
    ci_tests_run: int
    provenance: dict[tuple[str, str], dict] = field(default_factory=dict)


class PCAlgorithm:
    """PC algorithm for causal structure discovery.

    Two phases:
    1. Skeleton: remove edges where conditional independence holds
    2. Orientation: apply Meek's rules to orient remaining edges
    """

    def __init__(
        self,
        ci_test: PartialCorrelationTest | None = None,
        alpha: float = 0.05,
        max_cond_size: int = 3,
    ) -> None:
        self._ci_test = ci_test or PartialCorrelationTest()
        self._alpha = alpha
        self._max_cond_size = max_cond_size

    def fit(
        self,
        data: list[dict],
        variables: list[str],
        forbidden_edges: set[tuple[str, str]] | None = None,
        required_edges: set[tuple[str, str]] | None = None,
    ) -> PCResult:
        """Run PC algorithm on data."""
        forbidden = forbidden_edges or set()
        required = required_edges or set()

        # Phase 1: Skeleton discovery
        adjacency, sep_sets, ci_count, provenance = self._skeleton(
            data, variables, forbidden, required
        )

        # Phase 2: Orient edges
        oriented = self._orient(variables, adjacency, sep_sets, required)

        # Build skeleton as frozenset of sorted tuples
        skeleton_edges: set[tuple[str, str]] = set()
        for v in variables:
            for neighbor in adjacency.get(v, set()):
                edge = tuple(sorted([v, neighbor]))
                skeleton_edges.add((edge[0], edge[1]))

        return PCResult(
            skeleton=frozenset(skeleton_edges),
            oriented=frozenset(oriented),
            separating_sets=sep_sets,
            ci_tests_run=ci_count,
            provenance=provenance,
        )

    def _skeleton(
        self,
        data: list[dict],
        variables: list[str],
        forbidden: set[tuple[str, str]],
        required: set[tuple[str, str]],
    ) -> tuple[
        dict[str, set[str]],
        dict[tuple[str, str], frozenset[str]],
        int,
        dict[tuple[str, str], dict],
    ]:
        """Phase 1: Learn undirected skeleton.

        Optimized: pre-extracts all column data once and uses memoized
        partial correlations to avoid redundant computation.
        """
        # Start with complete graph (minus forbidden)
        adjacency: dict[str, set[str]] = {v: set() for v in variables}
        for i, v1 in enumerate(variables):
            for v2 in variables[i + 1 :]:
                edge_fwd = (v1, v2)
                edge_rev = (v2, v1)
                if edge_fwd not in forbidden and edge_rev not in forbidden:
                    adjacency[v1].add(v2)
                    adjacency[v2].add(v1)

        # Pre-extract ALL columns once (avoids re-parsing per CI test)
        n = len(data)
        col_vals: dict[str, list[float]] = {v: [] for v in variables}
        for row in data:
            for v in variables:
                val = row.get(v)
                if val is None:
                    col_vals[v].append(0.0)
                else:
                    try:
                        col_vals[v].append(float(val))
                    except (ValueError, TypeError):
                        col_vals[v].append(0.0)

        # Shared memoization cache for partial correlations across all tests
        corr_cache: dict[tuple[str, str, frozenset[str]], float] = {}

        sep_sets: dict[tuple[str, str], frozenset[str]] = {}
        provenance: dict[tuple[str, str], dict] = {}
        ci_count = 0

        for d in range(self._max_cond_size + 1):
            edges_to_check = []
            for v in variables:
                for neighbor in list(adjacency.get(v, set())):
                    if v < neighbor:
                        edges_to_check.append((v, neighbor))

            for x, y in edges_to_check:
                if y not in adjacency.get(x, set()):  # pragma: no cover
                    continue

                if (x, y) in required or (y, x) in required:
                    continue

                adj_x = adjacency[x] - {y}
                adj_y = adjacency[y] - {x}
                potential = adj_x | adj_y

                if len(potential) < d:
                    continue

                separated = False
                for z_set_tuple in combinations(sorted(potential), d):
                    z_set = frozenset(z_set_tuple)
                    result = self._ci_test.test_precomputed(
                        col_vals,
                        n,
                        x,
                        y,
                        z_set,
                        self._alpha,
                        corr_cache,
                    )
                    ci_count += 1

                    if result.independent:
                        adjacency[x].discard(y)
                        adjacency[y].discard(x)
                        key = tuple(sorted([x, y]))
                        sep_sets[(key[0], key[1])] = z_set
                        provenance[(key[0], key[1])] = {
                            "p_value": result.p_value,
                            "separating_set": sorted(z_set),
                            "algorithm": "pc",
                        }
                        separated = True
                        break

                if separated:
                    continue

        return adjacency, sep_sets, ci_count, provenance

    def _orient(
        self,
        variables: list[str],
        adjacency: dict[str, set[str]],
        sep_sets: dict[tuple[str, str], frozenset[str]],
        required: set[tuple[str, str]],
    ) -> set[tuple[str, str]]:
        """Phase 2: Orient edges using Meek's rules."""
        oriented: set[tuple[str, str]] = set()

        # Add required edges as already oriented
        for src, tgt in required:
            if tgt in adjacency.get(src, set()):
                oriented.add((src, tgt))

        # Rule 1: Colliders
        # If X—Z—Y and Z not in sep(X,Y) and X,Y not adjacent → X→Z←Y
        for z in variables:
            neighbors = sorted(adjacency.get(z, set()))
            for i, x in enumerate(neighbors):
                for y in neighbors[i + 1 :]:
                    # Check X and Y not adjacent
                    if y in adjacency.get(x, set()):
                        continue
                    # Check Z not in separating set of (X, Y)
                    key = tuple(sorted([x, y]))
                    sep = sep_sets.get((key[0], key[1]), frozenset())
                    if z not in sep:
                        oriented.add((x, z))
                        oriented.add((y, z))

        # Rules 2 & 3: iterate until no new orientations
        changed = True
        while changed:
            changed = False

            for z in variables:
                for y in adjacency.get(z, set()):
                    # Already oriented?
                    if (z, y) in oriented or (y, z) in oriented:
                        continue

                    # Rule 2: X→Z—Y and X,Y not adjacent → Z→Y
                    for x in variables:
                        if (x, z) in oriented and y in adjacency.get(z, set()):
                            if y not in adjacency.get(x, set()):
                                oriented.add((z, y))
                                changed = True
                                break

                    if (z, y) in oriented:
                        continue

                    # Rule 3: X→W, W→Z, X—Z → X→Z
                    for x in adjacency.get(z, set()):
                        if (x, z) in oriented or (z, x) in oriented:
                            continue
                        # Find W such that X→W and W→Z
                        for w in variables:
                            if (x, w) in oriented and (w, z) in oriented:
                                oriented.add((x, z))
                                changed = True
                                break
                        if (x, z) in oriented:
                            break

        return oriented
