"""Multi-algorithm consensus engine for causal discovery.

Runs PC and GES, keeps edges agreed by ≥ threshold fraction
of algorithms. Optionally scores edges via bootstrap stability.
Stdlib-only (ADR-007, ADR-017).
"""

from __future__ import annotations

import random
from dataclasses import dataclass

from orxaq.intelligence.ges_algorithm import GESAlgorithm
from orxaq.intelligence.independence import PartialCorrelationTest
from orxaq.intelligence.pc_algorithm import PCAlgorithm


@dataclass(frozen=True)
class ConsensusEdge:
    """A consensus edge with provenance from multiple algorithms."""

    source: str
    target: str
    directed: bool
    support_count: int
    support_fraction: float
    supporters: tuple[str, ...]
    confidence: str  # "high", "medium", "low"
    soliton_score: float | None = None


@dataclass(frozen=True)
class ConsensusResult:
    """Result of multi-algorithm consensus discovery."""

    edges: tuple[ConsensusEdge, ...]
    algorithms_used: tuple[str, ...]
    agreement_threshold: float


class ConsensusEngine:
    """Multi-algorithm consensus for causal discovery.

    Runs PC and GES, includes edges where ≥ threshold fraction agree.
    This is Orxaq's primary entry point for causal discovery.
    """

    def __init__(
        self,
        alpha: float = 0.05,
        threshold: float = 0.66,
        bootstrap_samples: int = 0,
        seed: int = 42,
    ) -> None:
        self._alpha = alpha
        self._threshold = threshold
        self._bootstrap_samples = bootstrap_samples
        self._seed = seed

    def discover(
        self,
        data: list[dict],
        variables: list[str],
        forbidden_edges: set[tuple[str, str]] | None = None,
        required_edges: set[tuple[str, str]] | None = None,
    ) -> ConsensusResult:
        """Run all algorithms and return consensus edges."""
        forbidden = forbidden_edges or set()
        required = required_edges or set()

        # Run PC algorithm
        ci_test = PartialCorrelationTest()
        pc = PCAlgorithm(ci_test=ci_test, alpha=self._alpha)
        pc_result = pc.fit(data, variables, forbidden, required)

        # Run GES algorithm
        ges = GESAlgorithm()
        ges_result = ges.fit(data, variables, forbidden, required)

        # Collect edges from each algorithm
        # PC: oriented edges are directed, skeleton edges are undirected
        pc_directed: set[tuple[str, str]] = set(pc_result.oriented)
        pc_undirected: set[tuple[str, str]] = set()
        for a, b in pc_result.skeleton:
            if (a, b) not in pc_directed and (b, a) not in pc_directed:
                pc_undirected.add((a, b))

        ges_directed: set[tuple[str, str]] = set(ges_result.edges)

        # Build edge support map
        all_edges: dict[tuple[str, str], list[str]] = {}
        edge_directed: dict[tuple[str, str], bool] = {}

        # Add PC directed edges
        for edge in pc_directed:
            key = edge
            all_edges.setdefault(key, []).append("pc")
            edge_directed[key] = True

        # Add PC undirected edges (normalize to sorted tuple)
        for a, b in pc_undirected:
            key = (min(a, b), max(a, b))
            all_edges.setdefault(key, []).append("pc")
            edge_directed.setdefault(key, False)

        # Add GES directed edges
        for edge in ges_directed:
            # Check if this matches an existing key in either direction
            key = edge
            reverse = (edge[1], edge[0])
            if reverse in all_edges and key not in all_edges:
                # GES disagrees on direction — add as undirected
                normalized = (min(edge[0], edge[1]), max(edge[0], edge[1]))
                # Move existing entry to normalized key
                if reverse in all_edges:
                    existing = all_edges.pop(reverse)
                    all_edges.setdefault(normalized, []).extend(existing)
                    edge_directed[normalized] = False
                all_edges.setdefault(normalized, []).append("ges")
                edge_directed[normalized] = False
            else:
                all_edges.setdefault(key, []).append("ges")
                edge_directed.setdefault(key, True)

        # Filter by threshold
        n_algorithms = 2
        consensus_edges: list[ConsensusEdge] = []

        for edge, supporters in sorted(all_edges.items()):
            unique_supporters = sorted(set(supporters))
            count = len(unique_supporters)
            fraction = count / n_algorithms

            if fraction < self._threshold:
                continue

            # Determine confidence
            if count == n_algorithms:
                confidence = "high"
            elif fraction >= self._threshold:
                confidence = "medium"
            else:  # pragma: no cover
                confidence = "low"  # pragma: no cover

            # Compute soliton score if bootstrap requested
            soliton = None
            if self._bootstrap_samples > 0:
                soliton = self._bootstrap_stability(data, variables, edge, forbidden, required)

            consensus_edges.append(
                ConsensusEdge(
                    source=edge[0],
                    target=edge[1],
                    directed=edge_directed.get(edge, True),
                    support_count=count,
                    support_fraction=fraction,
                    supporters=tuple(unique_supporters),
                    confidence=confidence,
                    soliton_score=soliton,
                )
            )

        return ConsensusResult(
            edges=tuple(consensus_edges),
            algorithms_used=("pc", "ges"),
            agreement_threshold=self._threshold,
        )

    def _bootstrap_stability(
        self,
        data: list[dict],
        variables: list[str],
        edge: tuple[str, str],
        forbidden: set[tuple[str, str]],
        required: set[tuple[str, str]],
    ) -> float:
        """Compute bootstrap stability for an edge.

        Run discovery on bootstrap samples and count how often the edge appears.
        """
        rng = random.Random(self._seed)
        n = len(data)
        appearances = 0

        for _ in range(self._bootstrap_samples):
            # Bootstrap resample
            sample = [data[rng.randint(0, n - 1)] for _ in range(n)]

            # Run quick PC on bootstrap sample
            ci_test = PartialCorrelationTest()
            pc = PCAlgorithm(ci_test=ci_test, alpha=self._alpha, max_cond_size=1)
            result = pc.fit(sample, variables, forbidden, required)

            # Check if edge appears in skeleton or oriented
            edge_found = False
            for a, b in result.skeleton:
                if (a, b) == edge or (b, a) == edge:
                    edge_found = True
                    break
            if not edge_found:
                for a, b in result.oriented:
                    if (a, b) == edge or (b, a) == edge:
                        edge_found = True
                        break

            if edge_found:
                appearances += 1

        return appearances / self._bootstrap_samples if self._bootstrap_samples > 0 else 0.0
