"""GES algorithm — score-based causal discovery.

Greedy Equivalence Search: discovers causal structure by optimizing BIC.
Forward phase adds edges, backward phase removes them.
Stdlib-only (ADR-007).
"""

from __future__ import annotations

import math
from collections import deque
from dataclasses import dataclass


@dataclass(frozen=True)
class GESResult:
    """Result of GES algorithm discovery."""

    edges: frozenset[tuple[str, str]]
    bic_score: float
    additions: int
    removals: int


class GESAlgorithm:
    """Greedy Equivalence Search for causal structure discovery.

    Two phases:
    1. Forward: greedily add edges that improve BIC
    2. Backward: greedily remove edges that improve BIC
    """

    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,
    ) -> GESResult:
        """Run GES algorithm on data."""
        forbidden = forbidden_edges or set()
        required = required_edges or set()

        n = len(data)
        if n < 5 or len(variables) < 2:
            return GESResult(edges=frozenset(required), bic_score=0.0, additions=0, removals=0)

        # Extract numeric data once upfront
        col_vals = self._extract_columns(data, variables)

        # Initialize with required edges + parent map for O(1) lookups
        edges: set[tuple[str, str]] = set(required)
        parent_map: dict[str, list[str]] = {v: [] for v in variables}
        children_map: dict[str, set[str]] = {v: set() for v in variables}
        for src, tgt in required:
            parent_map[tgt].append(src)
            children_map.setdefault(src, set()).add(tgt)

        # Pre-compute log(n) once
        log_n = math.log(n)

        # Forward phase: add edges that improve BIC
        additions = 0
        while True:
            best_edge = None
            best_improvement = 0.0

            for i, src in enumerate(variables):
                for tgt in variables[i + 1 :]:
                    for edge in [(src, tgt), (tgt, src)]:
                        if edge in edges or edge in forbidden:
                            continue
                        reverse = (edge[1], edge[0])
                        if reverse in edges or reverse in forbidden:
                            continue

                        if self._would_cycle(edge[0], edge[1], children_map):
                            continue

                        improvement = self._bic_improvement_add(
                            col_vals,
                            edge,
                            parent_map,
                            n,
                            log_n,
                        )
                        if improvement > best_improvement:
                            best_improvement = improvement
                            best_edge = edge

            if best_edge is None:
                break
            edges.add(best_edge)
            parent_map[best_edge[1]].append(best_edge[0])
            children_map.setdefault(best_edge[0], set()).add(best_edge[1])
            additions += 1

        # Backward phase: remove edges that improve BIC
        removals = 0
        while True:
            best_edge = None
            best_improvement = 0.0

            for edge in list(edges):
                if edge in required:
                    continue
                improvement = self._bic_improvement_remove(
                    col_vals,
                    edge,
                    parent_map,
                    n,
                    log_n,
                )
                if improvement > best_improvement:
                    best_improvement = improvement
                    best_edge = edge

            if best_edge is None:
                break
            edges.discard(best_edge)
            parent_map[best_edge[1]] = [p for p in parent_map[best_edge[1]] if p != best_edge[0]]
            children_map.get(best_edge[0], set()).discard(best_edge[1])
            removals += 1

        bic = self._total_bic(col_vals, parent_map, variables, n, log_n)

        return GESResult(
            edges=frozenset(edges),
            bic_score=bic,
            additions=additions,
            removals=removals,
        )

    @staticmethod
    def _extract_columns(data: list[dict], variables: list[str]) -> dict[str, list[float]]:
        """Extract column values as floats."""
        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)
        return col_vals

    @staticmethod
    def _would_cycle(
        source: str,
        target: str,
        children_map: dict[str, set[str]],
    ) -> bool:
        """Check if adding source->target would create a cycle.

        Uses BFS with deque + adjacency dict for O(V+E) instead of O(V*E).
        """
        visited: set[str] = set()
        queue: deque[str] = deque([target])
        while queue:
            node = queue.popleft()
            if node == source:
                return True
            if node in visited:
                continue
            visited.add(node)
            queue.extend(children_map.get(node, set()) - visited)
        return False

    def _bic_improvement_add(
        self,
        col_vals: dict[str, list[float]],
        edge: tuple[str, str],
        parent_map: dict[str, list[str]],
        n: int,
        log_n: float,
    ) -> float:
        """BIC improvement from adding an edge."""
        target = edge[1]
        current_parents = parent_map[target]
        new_parents = current_parents + [edge[0]]

        bic_before = self._node_bic(col_vals, target, current_parents, n, log_n)
        bic_after = self._node_bic(col_vals, target, new_parents, n, log_n)

        return bic_before - bic_after

    def _bic_improvement_remove(
        self,
        col_vals: dict[str, list[float]],
        edge: tuple[str, str],
        parent_map: dict[str, list[str]],
        n: int,
        log_n: float,
    ) -> float:
        """BIC improvement from removing an edge."""
        target = edge[1]
        current_parents = parent_map[target]
        new_parents = [p for p in current_parents if p != edge[0]]

        bic_before = self._node_bic(col_vals, target, current_parents, n, log_n)
        bic_after = self._node_bic(col_vals, target, new_parents, n, log_n)

        return bic_before - bic_after

    @staticmethod
    def _node_bic(
        col_vals: dict[str, list[float]],
        target: str,
        parents: list[str],
        n: int,
        log_n: float,
    ) -> float:
        """Compute BIC score for a single node given its parents.

        BIC = n * ln(RSS/n) + |parents| * ln(n)
        """
        y = col_vals[target]

        if not parents:
            mean_y = sum(y) / n
            rss = sum((yi - mean_y) ** 2 for yi in y)
        else:
            rss = _ols_rss(col_vals, target, parents, n)

        rss = max(rss, 1e-15)
        return n * math.log(rss / n) + len(parents) * log_n

    def _total_bic(
        self,
        col_vals: dict[str, list[float]],
        parent_map: dict[str, list[str]],
        variables: list[str],
        n: int,
        log_n: float,
    ) -> float:
        """Total BIC score across all nodes."""
        total = 0.0
        for v in variables:
            total += self._node_bic(col_vals, v, parent_map.get(v, []), n, log_n)
        return total


def _ols_rss(
    col_vals: dict[str, list[float]],
    target: str,
    parents: list[str],
    n: int,
) -> float:
    """Compute residual sum of squares from OLS regression.

    Uses normal equations: beta = (X'X)^{-1} X'y
    For small parent sets, direct matrix inversion via Cramer's rule or
    iterative approach.
    """
    y = col_vals[target]
    k = len(parents)

    if k == 0:
        mean_y = sum(y) / n
        return sum((yi - mean_y) ** 2 for yi in y)

    # Build X matrix (with intercept)
    # Compute X'X and X'y directly
    p = k + 1  # +1 for intercept

    # X'X matrix (p x p)
    xtx = [[0.0] * p for _ in range(p)]
    xty = [0.0] * p

    for i in range(n):
        row = [1.0] + [col_vals[par][i] for par in parents]
        for j in range(p):
            xty[j] += row[j] * y[i]
            for m in range(p):
                xtx[j][m] += row[j] * row[m]

    # Solve via Gaussian elimination
    beta = _solve_linear(xtx, xty, p)
    if beta is None:
        # Singular matrix — return total SS
        mean_y = sum(y) / n
        return sum((yi - mean_y) ** 2 for yi in y)

    # Compute RSS
    rss = 0.0
    for i in range(n):
        row = [1.0] + [col_vals[par][i] for par in parents]
        predicted = sum(beta[j] * row[j] for j in range(p))
        rss += (y[i] - predicted) ** 2

    return rss


def _solve_linear(a: list[list[float]], b: list[float], n: int) -> list[float] | None:
    """Solve Ax = b via Gaussian elimination with partial pivoting."""
    # Augmented matrix
    aug = [a[i][:] + [b[i]] for i in range(n)]

    for col in range(n):
        # Partial pivoting
        max_row = col
        for row in range(col + 1, n):
            if abs(aug[row][col]) > abs(aug[max_row][col]):
                max_row = row
        aug[col], aug[max_row] = aug[max_row], aug[col]

        if abs(aug[col][col]) < 1e-12:
            return None

        # Eliminate below
        for row in range(col + 1, n):
            factor = aug[row][col] / aug[col][col]
            for j in range(col, n + 1):
                aug[row][j] -= factor * aug[col][j]

    # Back substitution
    x = [0.0] * n
    for i in range(n - 1, -1, -1):
        x[i] = aug[i][n]
        for j in range(i + 1, n):
            x[i] -= aug[i][j] * x[j]
        x[i] /= aug[i][i]

    return x
