"""Conditional independence tests for causal discovery.

Foundation for constraint-based algorithms: tests whether X is independent
of Y given a conditioning set Z, using partial correlation and Fisher's
z-transform. All stdlib-only (ADR-007).
"""

from __future__ import annotations

import math
from dataclasses import dataclass


@dataclass(frozen=True)
class CIResult:
    """Result of a conditional independence test."""

    x: str
    y: str
    z_set: frozenset[str]
    statistic: float
    p_value: float
    independent: bool
    method: str


def pearson_correlation(x_vals: list[float], y_vals: list[float]) -> float:
    """Compute Pearson correlation coefficient (stdlib only)."""
    n = len(x_vals)
    if n < 3:
        return 0.0

    mean_x = sum(x_vals) / n
    mean_y = sum(y_vals) / n

    cov_xy = sum((x - mean_x) * (y - mean_y) for x, y in zip(x_vals, y_vals)) / n
    var_x = sum((x - mean_x) ** 2 for x in x_vals) / n
    var_y = sum((y - mean_y) ** 2 for y in y_vals) / n

    denom = math.sqrt(var_x * var_y)
    if denom < 1e-15:
        return 0.0

    return max(-1.0, min(1.0, cov_xy / denom))


def _normal_cdf(z: float) -> float:
    """Standard normal CDF approximation via math.erf."""
    return 0.5 * (1.0 + math.erf(z / math.sqrt(2.0)))


class PartialCorrelationTest:
    """Conditional independence test via partial correlation.

    Uses recursive partial correlation formula and Fisher's z-transform
    for significance testing. Supports pre-extracted columns via
    test_precomputed() to avoid repeated data parsing in hot loops.
    """

    def test(
        self,
        data: list[dict],
        x: str,
        y: str,
        z_set: frozenset[str] | None = None,
        alpha: float = 0.05,
    ) -> CIResult:
        """Test if X is conditionally independent of Y given Z."""
        if z_set is None:
            z_set = frozenset()

        n = len(data)
        if n < 5:
            return CIResult(
                x=x,
                y=y,
                z_set=z_set,
                statistic=0.0,
                p_value=1.0,
                independent=True,
                method="partial_correlation",
            )

        # Extract column values
        columns = {x, y} | z_set
        col_vals: dict[str, list[float]] = {c: [] for c in columns}
        for row in data:
            for c in columns:
                val = row.get(c)
                if val is None:
                    col_vals[c].append(0.0)
                else:
                    try:
                        col_vals[c].append(float(val))
                    except (ValueError, TypeError):
                        col_vals[c].append(0.0)

        r = self._partial_corr(col_vals, x, y, z_set)
        p_value = self._fisher_z_pvalue(r, n, len(z_set))

        return CIResult(
            x=x,
            y=y,
            z_set=z_set,
            statistic=r,
            p_value=p_value,
            independent=p_value > alpha,
            method="partial_correlation",
        )

    def test_precomputed(
        self,
        col_vals: dict[str, list[float]],
        n: int,
        x: str,
        y: str,
        z_set: frozenset[str],
        alpha: float = 0.05,
        _cache: dict[tuple[str, str, frozenset[str]], float] | None = None,
    ) -> CIResult:
        """Test CI with pre-extracted column data (avoids repeated parsing)."""
        if n < 5:
            return CIResult(
                x=x,
                y=y,
                z_set=z_set,
                statistic=0.0,
                p_value=1.0,
                independent=True,
                method="partial_correlation",
            )

        r = self._partial_corr(col_vals, x, y, z_set, _cache)
        p_value = self._fisher_z_pvalue(r, n, len(z_set))

        return CIResult(
            x=x,
            y=y,
            z_set=z_set,
            statistic=r,
            p_value=p_value,
            independent=p_value > alpha,
            method="partial_correlation",
        )

    def _partial_corr(
        self,
        col_vals: dict[str, list[float]],
        x: str,
        y: str,
        z_set: frozenset[str],
        _cache: dict[tuple[str, str, frozenset[str]], float] | None = None,
    ) -> float:
        """Recursive partial correlation with optional memoization.

        Memoization reduces O(3^k) recursive calls to O(k^2) unique computations.
        """
        if _cache is not None:
            key = (x, y, z_set) if x <= y else (y, x, z_set)
            if key in _cache:
                return _cache[key]

        if not z_set:
            result = pearson_correlation(col_vals[x], col_vals[y])
            if _cache is not None:
                key = (x, y, z_set) if x <= y else (y, x, z_set)
                _cache[key] = result
            return result

        z_list = sorted(z_set)
        z = z_list[-1]
        remaining = frozenset(z_list[:-1])

        r_xy = self._partial_corr(col_vals, x, y, remaining, _cache)
        r_xz = self._partial_corr(col_vals, x, z, remaining, _cache)
        r_yz = self._partial_corr(col_vals, y, z, remaining, _cache)

        denom = math.sqrt(max(1e-15, (1 - r_xz**2) * (1 - r_yz**2)))
        result = max(-1.0, min(1.0, (r_xy - r_xz * r_yz) / denom))

        if _cache is not None:
            key = (x, y, z_set) if x <= y else (y, x, z_set)
            _cache[key] = result
        return result

    @staticmethod
    def _fisher_z_pvalue(r: float, n: int, k: int) -> float:
        """Fisher z-transform p-value.

        z = 0.5 * ln((1+r)/(1-r))
        Under H0: z ~ N(0, 1/sqrt(n - |Z| - 3))
        """
        # Clamp r to avoid log(0)
        r = max(-0.9999, min(0.9999, r))
        z = 0.5 * math.log((1 + r) / (1 - r))

        df = n - k - 3
        if df < 1:
            return 1.0

        se = 1.0 / math.sqrt(df)
        z_stat = abs(z) / se

        # Two-tailed p-value
        return 2.0 * (1.0 - _normal_cdf(z_stat))
