"""Formal do-calculus intervention analysis for causal queries.

Implements Pearl's do-calculus framework:
- Backdoor criterion & adjustment sets
- Frontdoor criterion & identification
- Instrumental variable identification
- Causal effect estimation (IPW, stratification)
- Sensitivity analysis (Rosenbaum bounds)
- Mediation analysis (direct/indirect effects)

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

from __future__ import annotations

import math
from collections import deque
from dataclasses import dataclass
from enum import Enum

from orxaq.kernel import WorldDAG


class IdentificationMethod(Enum):
    """Strategy used to identify a causal effect."""

    BACKDOOR = "backdoor"
    FRONTDOOR = "frontdoor"
    INSTRUMENTAL_VARIABLE = "instrumental_variable"
    NONE = "none"


@dataclass(frozen=True)
class AdjustmentSet:
    """A valid adjustment set for causal identification."""

    variables: frozenset[str]
    method: IdentificationMethod
    is_minimal: bool = False


@dataclass(frozen=True)
class IdentificationResult:
    """Result of checking whether a causal effect is identifiable."""

    identifiable: bool
    method: IdentificationMethod
    adjustment_sets: tuple[AdjustmentSet, ...]
    instruments: tuple[str, ...] = ()


@dataclass(frozen=True)
class InterventionResult:
    """Estimated causal effect from an intervention."""

    ate: float
    ci_lower: float
    ci_upper: float
    method: IdentificationMethod
    adjustment_set: frozenset[str]
    n_samples: int


@dataclass(frozen=True)
class SensitivityResult:
    """Sensitivity analysis for unmeasured confounding."""

    gamma: float
    p_value_upper: float
    p_value_lower: float
    robust: bool


@dataclass(frozen=True)
class MediationResult:
    """Mediation analysis decomposition."""

    total_effect: float
    direct_effect: float
    indirect_effect: float
    proportion_mediated: float
    mediator: str


# ── Helpers (stdlib-only stats) ────────────────────────────────


def _covariance(xs: list[float], ys: list[float]) -> float:
    """Sample covariance between two sequences."""
    n = len(xs)
    if n < 2:
        return 0.0
    mx = sum(xs) / n
    my = sum(ys) / n
    return sum((x - mx) * (y - my) for x, y in zip(xs, ys)) / (n - 1)


def _variance(xs: list[float]) -> float:
    """Sample variance of a sequence."""
    n = len(xs)
    if n < 2:
        return 0.0
    mx = sum(xs) / n
    return sum((x - mx) ** 2 for x in xs) / (n - 1)


def _regression_coefficient(xs: list[float], ys: list[float]) -> float:
    """OLS regression coefficient of Y on X."""
    var_x = _variance(xs)
    if var_x < 1e-15:
        return 0.0
    return _covariance(xs, ys) / var_x


def _partial_regression_coefficient(
    xs: list[float],
    ys: list[float],
    zs: list[list[float]],
) -> float:
    """Partial regression coefficient of Y on X controlling for Z.

    Uses Frisch-Waugh-Lovell: regress X on Z, regress Y on Z,
    then regress residuals.
    """
    if not zs:
        return _regression_coefficient(xs, ys)

    # Residualise X and Y w.r.t. each Z variable sequentially
    rx = list(xs)
    ry = list(ys)
    for z_col in zs:
        bx = _regression_coefficient(z_col, rx)
        by = _regression_coefficient(z_col, ry)
        rx = [r - bx * z for r, z in zip(rx, z_col)]
        ry = [r - by * z for r, z in zip(ry, z_col)]

    return _regression_coefficient(rx, ry)


def _bin_index(value: float, n_bins: int, lo: float, hi: float) -> int:
    """Assign a value to a bin index in [0, n_bins)."""
    if hi <= lo:
        return 0
    idx = int((value - lo) / (hi - lo) * n_bins)
    return max(0, min(n_bins - 1, idx))


def _combinations(items: list[str], r: int) -> list[frozenset[str]]:
    """Generate all r-combinations of items as frozensets."""
    result: list[frozenset[str]] = []
    n = len(items)
    if r > n or r < 0:
        return result
    indices = list(range(r))
    result.append(frozenset(items[i] for i in indices))
    while True:
        pivot = -1
        for i in range(r - 1, -1, -1):
            if indices[i] != i + n - r:
                pivot = i
                break
        if pivot == -1:
            break
        indices[pivot] += 1
        for j in range(pivot + 1, r):
            indices[j] = indices[j - 1] + 1
        result.append(frozenset(items[i] for i in indices))
    return result


# ── D-separation in mutilated graphs ──────────────────────────


def _dsep_in_mutilated_graph(
    dag: WorldDAG,
    source: str,
    target: str,
    z: frozenset[str],
    mutilated_node: str | None = None,
) -> bool:
    """Test d-separation in graph where outgoing edges of mutilated_node removed.

    This implements the do(X) operation: G_{overline{X}} removes all
    edges OUT of X, simulating an intervention that breaks incoming
    causal influences on X.
    """
    conditioned_set = set(z)

    # Precompute ancestors of conditioned set (in mutilated graph)
    has_cond_desc: set[str] = set()
    for c in conditioned_set:
        has_cond_desc.add(c)
        # Walk up parents (respecting mutation)
        queue: deque[str] = deque()
        for p in dag._parents_map.get(c, set()):
            queue.append(p)
        while queue:
            node = queue.popleft()
            if node in has_cond_desc:
                continue
            has_cond_desc.add(node)
            for p in dag._parents_map.get(node, set()):
                queue.append(p)

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

    bb_queue: deque[tuple[str, bool]] = deque()

    # Initialize from source
    for parent in dag._parents_map.get(source, set()):
        bb_queue.append((parent, True))
    # Only send down if source is not mutilated
    if source != mutilated_node:
        for child in dag._children_map.get(source, set()):
            bb_queue.append((child, False))

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

        if from_child and node not in conditioned_set:
            if node not in visited_from_child:
                visited_from_child.add(node)
                reachable.add(node)
                for parent in dag._parents_map.get(node, set()):
                    bb_queue.append((parent, True))
                # Only propagate down if not mutilated
                if node != mutilated_node:
                    for child in dag._children_map.get(node, set()):
                        bb_queue.append((child, False))

        elif from_child and node in conditioned_set:
            pass  # Blocked by conditioning

        elif not from_child and node not in conditioned_set:
            if node not in visited_from_parent:
                visited_from_parent.add(node)
                reachable.add(node)
                if node != mutilated_node:
                    for child in dag._children_map.get(node, set()):
                        bb_queue.append((child, False))

        elif not from_child and node in conditioned_set:
            if node not in visited_from_parent:
                visited_from_parent.add(node)
                for parent in dag._parents_map.get(node, set()):
                    bb_queue.append((parent, True))

    reachable.discard(source)
    return target not in reachable


# ── Adjustment Set Finder ─────────────────────────────────────


class AdjustmentSetFinder:
    """Find valid adjustment sets for backdoor identification.

    A set Z satisfies the backdoor criterion relative to (X, Y) if:
    1. No node in Z is a descendant of X
    2. Z blocks every path between X and Y that contains an arrow into X
       (i.e., X ⊥ Y | Z in the mutilated graph G_{overline{X}})
    """

    def __init__(self, dag: WorldDAG) -> None:
        self._dag = dag

    def find_all(
        self,
        treatment: str,
        outcome: str,
        max_size: int = 3,
    ) -> list[AdjustmentSet]:
        """Find all valid backdoor adjustment sets up to max_size."""
        descendants_of_treatment = self._dag.descendants(treatment)
        candidates = [
            n
            for n in sorted(self._dag._entities)
            if n != treatment and n != outcome and n not in descendants_of_treatment
        ]

        valid: list[AdjustmentSet] = []
        for size in range(max_size + 1):
            for combo in _combinations(candidates, size):
                if self._satisfies_backdoor(treatment, outcome, combo):
                    valid.append(
                        AdjustmentSet(
                            variables=combo,
                            method=IdentificationMethod.BACKDOOR,
                            is_minimal=(size == 0),
                        )
                    )

        # Mark truly minimal sets
        if valid:
            min_size = min(len(a.variables) for a in valid)
            valid = [
                AdjustmentSet(
                    variables=a.variables,
                    method=a.method,
                    is_minimal=(len(a.variables) == min_size),
                )
                for a in valid
            ]

        return valid

    def find_minimal(self, treatment: str, outcome: str) -> AdjustmentSet | None:
        """Find the smallest valid backdoor adjustment set."""
        sets = self.find_all(treatment, outcome)
        if not sets:
            return None
        return min(sets, key=lambda a: len(a.variables))

    def _satisfies_backdoor(
        self,
        treatment: str,
        outcome: str,
        z: frozenset[str],
    ) -> bool:
        """Check if Z satisfies the backdoor criterion."""
        return _dsep_in_mutilated_graph(self._dag, treatment, outcome, z, mutilated_node=treatment)


# ── Identifiability Checker ───────────────────────────────────


class IdentifiabilityChecker:
    """Check whether a causal effect is identifiable.

    Tries in order: backdoor, frontdoor, instrumental variable.
    """

    def __init__(self, dag: WorldDAG) -> None:
        self._dag = dag
        self._adj_finder = AdjustmentSetFinder(dag)

    def check(self, treatment: str, outcome: str) -> IdentificationResult:
        """Check identifiability of P(Y | do(X))."""
        # Try backdoor
        adj_sets = self._adj_finder.find_all(treatment, outcome)
        if adj_sets:
            return IdentificationResult(
                identifiable=True,
                method=IdentificationMethod.BACKDOOR,
                adjustment_sets=tuple(adj_sets),
            )

        # Try frontdoor
        fd_sets = self._find_frontdoor_sets(treatment, outcome)
        if fd_sets:
            return IdentificationResult(
                identifiable=True,
                method=IdentificationMethod.FRONTDOOR,
                adjustment_sets=tuple(fd_sets),
            )

        # Try instrumental variable
        instruments = self._find_instruments(treatment, outcome)
        if instruments:
            return IdentificationResult(
                identifiable=True,
                method=IdentificationMethod.INSTRUMENTAL_VARIABLE,
                adjustment_sets=(),
                instruments=tuple(sorted(instruments)),
            )

        return IdentificationResult(
            identifiable=False,
            method=IdentificationMethod.NONE,
            adjustment_sets=(),
        )

    def _find_frontdoor_sets(self, treatment: str, outcome: str) -> list[AdjustmentSet]:
        """Find sets satisfying the frontdoor criterion.

        M satisfies frontdoor relative to (X, Y) if:
        1. X blocks all directed paths from X to M (trivially true
           since X intercepts all X→...→M paths)
           Actually: M intercepts all directed paths from X to Y
        2. No unblocked backdoor path from X to M
        3. All backdoor paths from M to Y are blocked by X
        """
        desc_t = self._dag.descendants(treatment)
        anc_o = self._dag.ancestors(outcome)

        mediators = [
            n
            for n in sorted(self._dag._entities)
            if n != treatment and n != outcome and n in desc_t and (n in anc_o or n == outcome)
        ]

        results: list[AdjustmentSet] = []
        for m in mediators:
            if self._is_frontdoor_mediator(treatment, outcome, m):
                results.append(
                    AdjustmentSet(
                        variables=frozenset({m}),
                        method=IdentificationMethod.FRONTDOOR,
                        is_minimal=True,
                    )
                )
        return results

    def _is_frontdoor_mediator(self, treatment: str, outcome: str, mediator: str) -> bool:
        """Check all three frontdoor conditions for a mediator."""
        # Condition 1: All directed paths from X to Y go through M
        paths = self._dag.all_directed_paths(treatment, outcome)
        if not paths:
            return False
        for path in paths:
            if mediator not in path:
                return False

        # Condition 2: No unblocked backdoor path from X to M
        # (X d-sep from M in G_{overline{X}})
        if not _dsep_in_mutilated_graph(
            self._dag,
            treatment,
            mediator,
            frozenset(),
            mutilated_node=treatment,
        ):
            return False

        # Condition 3: All backdoor paths from M to Y blocked by X
        x_set = frozenset({treatment})
        if not _dsep_in_mutilated_graph(
            self._dag,
            mediator,
            outcome,
            x_set,
            mutilated_node=mediator,
        ):
            return False

        return True

    def _find_instruments(self, treatment: str, outcome: str) -> list[str]:
        """Find instrumental variables.

        Z is an instrument for (X, Y) if:
        1. Z is associated with X (not d-separated from X)
        2. Z has no direct effect on Y (no directed path Z→Y
           that doesn't go through X)
        3. Z ⊥ Y in G_{overline{X}} (no confounding Z-Y path)
        """
        results: list[str] = []
        desc_t = self._dag.descendants(treatment)

        for z in sorted(self._dag._entities):
            if z == treatment or z == outcome:
                continue
            if z in desc_t:
                continue

            # Condition 1: Z associated with X
            if self._dag.d_separated(z, treatment, frozenset()):
                continue

            # Condition 2: All directed Z→Y paths go through X
            z_to_y_paths = self._dag.all_directed_paths(z, outcome)
            bypasses = False
            for path in z_to_y_paths:
                if treatment not in path:
                    bypasses = True
                    break
            if bypasses:
                continue

            # Condition 3: Z ⊥ Y in mutilated graph
            if not _dsep_in_mutilated_graph(
                self._dag,
                z,
                outcome,
                frozenset(),
                mutilated_node=treatment,
            ):
                continue

            results.append(z)

        return results


# ── Intervention Estimator ────────────────────────────────────


class InterventionEstimator:
    """Estimate causal effects using identified adjustment strategies."""

    def __init__(self, dag: WorldDAG) -> None:
        self._dag = dag
        self._checker = IdentifiabilityChecker(dag)

    def estimate_ate(
        self,
        data: list[dict[str, float]],
        treatment: str,
        outcome: str,
        method: IdentificationMethod | None = None,
    ) -> InterventionResult:
        """Estimate average treatment effect P(Y | do(X=1)) - P(Y | do(X=0)).

        Auto-selects identification strategy if method is None.
        """
        ident = self._checker.check(treatment, outcome)
        if not ident.identifiable:
            return InterventionResult(
                ate=0.0,
                ci_lower=0.0,
                ci_upper=0.0,
                method=IdentificationMethod.NONE,
                adjustment_set=frozenset(),
                n_samples=len(data),
            )

        use_method = method or ident.method

        if use_method == IdentificationMethod.BACKDOOR:
            adj = ident.adjustment_sets[0].variables
            return self._backdoor_estimate(data, treatment, outcome, adj)
        if use_method == IdentificationMethod.INSTRUMENTAL_VARIABLE:
            iv = ident.instruments[0]
            return self._iv_estimate(data, treatment, outcome, iv)
        if use_method == IdentificationMethod.FRONTDOOR:
            mediator = next(iter(ident.adjustment_sets[0].variables))
            return self._frontdoor_estimate(data, treatment, outcome, mediator)

        return InterventionResult(
            ate=0.0,
            ci_lower=0.0,
            ci_upper=0.0,
            method=IdentificationMethod.NONE,
            adjustment_set=frozenset(),
            n_samples=len(data),
        )

    def _backdoor_estimate(
        self,
        data: list[dict[str, float]],
        treatment: str,
        outcome: str,
        adjustment: frozenset[str],
    ) -> InterventionResult:
        """Estimate via backdoor adjustment (stratification)."""
        if not adjustment:
            return self._simple_difference(data, treatment, outcome, adjustment)

        adj_vars = sorted(adjustment)
        n_bins = 5

        # Get ranges for binning
        ranges: dict[str, tuple[float, float]] = {}
        for v in adj_vars:
            vals = [d[v] for d in data]
            ranges[v] = (min(vals), max(vals))

        # Stratify
        strata: dict[tuple[int, ...], list[dict[str, float]]] = {}
        for row in data:
            key = tuple(_bin_index(row[v], n_bins, *ranges[v]) for v in adj_vars)
            strata.setdefault(key, []).append(row)

        # Compute weighted ATE across strata
        total_n = len(data)
        ate = 0.0
        for stratum_data in strata.values():
            w = len(stratum_data) / total_n
            treated = [d[outcome] for d in stratum_data if d[treatment] >= 0.5]
            control = [d[outcome] for d in stratum_data if d[treatment] < 0.5]
            if treated and control:
                diff = sum(treated) / len(treated) - sum(control) / len(control)
                ate += w * diff

        se = self._bootstrap_se(data, treatment, outcome, adjustment)
        return InterventionResult(
            ate=ate,
            ci_lower=ate - 1.96 * se,
            ci_upper=ate + 1.96 * se,
            method=IdentificationMethod.BACKDOOR,
            adjustment_set=adjustment,
            n_samples=len(data),
        )

    def _simple_difference(
        self,
        data: list[dict[str, float]],
        treatment: str,
        outcome: str,
        adjustment: frozenset[str],
    ) -> InterventionResult:
        """Simple mean difference when no adjustment needed."""
        treated = [d[outcome] for d in data if d[treatment] >= 0.5]
        control = [d[outcome] for d in data if d[treatment] < 0.5]
        if not treated or not control:
            return InterventionResult(
                ate=0.0,
                ci_lower=0.0,
                ci_upper=0.0,
                method=IdentificationMethod.BACKDOOR,
                adjustment_set=adjustment,
                n_samples=len(data),
            )
        ate = sum(treated) / len(treated) - sum(control) / len(control)
        se = self._bootstrap_se(data, treatment, outcome, adjustment)
        return InterventionResult(
            ate=ate,
            ci_lower=ate - 1.96 * se,
            ci_upper=ate + 1.96 * se,
            method=IdentificationMethod.BACKDOOR,
            adjustment_set=adjustment,
            n_samples=len(data),
        )

    def _iv_estimate(
        self,
        data: list[dict[str, float]],
        treatment: str,
        outcome: str,
        instrument: str,
    ) -> InterventionResult:
        """Two-stage least squares IV estimate."""
        z = [d[instrument] for d in data]
        x = [d[treatment] for d in data]
        y = [d[outcome] for d in data]

        cov_zy = _covariance(z, y)
        cov_zx = _covariance(z, x)

        if abs(cov_zx) < 1e-15:
            return InterventionResult(
                ate=0.0,
                ci_lower=0.0,
                ci_upper=0.0,
                method=IdentificationMethod.INSTRUMENTAL_VARIABLE,
                adjustment_set=frozenset(),
                n_samples=len(data),
            )

        ate = cov_zy / cov_zx
        n = len(data)
        se = abs(ate) / math.sqrt(max(n - 2, 1)) if n > 2 else 0.0
        return InterventionResult(
            ate=ate,
            ci_lower=ate - 1.96 * se,
            ci_upper=ate + 1.96 * se,
            method=IdentificationMethod.INSTRUMENTAL_VARIABLE,
            adjustment_set=frozenset(),
            n_samples=n,
        )

    def _frontdoor_estimate(
        self,
        data: list[dict[str, float]],
        treatment: str,
        outcome: str,
        mediator: str,
    ) -> InterventionResult:
        """Frontdoor adjustment estimate.

        ATE = sum_m P(m|do(x=1)) * E[Y|do(M=m)]
            - sum_m P(m|do(x=0)) * E[Y|do(M=m)]
        Simplified via regression: beta_xm * beta_my_x
        """
        x = [d[treatment] for d in data]
        m = [d[mediator] for d in data]
        y = [d[outcome] for d in data]

        beta_xm = _regression_coefficient(x, m)
        beta_my = _partial_regression_coefficient(m, y, [x])
        ate = beta_xm * beta_my

        n = len(data)
        se = abs(ate) / math.sqrt(max(n - 2, 1)) if n > 2 else 0.0
        return InterventionResult(
            ate=ate,
            ci_lower=ate - 1.96 * se,
            ci_upper=ate + 1.96 * se,
            method=IdentificationMethod.FRONTDOOR,
            adjustment_set=frozenset({mediator}),
            n_samples=n,
        )

    def _bootstrap_se(
        self,
        data: list[dict[str, float]],
        treatment: str,
        outcome: str,
        adjustment: frozenset[str],
        n_bootstrap: int = 50,
        seed: int = 42,
    ) -> float:
        """Quick bootstrap standard error estimate."""
        import random as _rng

        gen = _rng.Random(seed)
        n = len(data)
        estimates: list[float] = []

        for _ in range(n_bootstrap):
            sample = [data[gen.randint(0, n - 1)] for _ in range(n)]
            treated = [d[outcome] for d in sample if d[treatment] >= 0.5]
            control = [d[outcome] for d in sample if d[treatment] < 0.5]
            if treated and control:
                est = sum(treated) / len(treated) - sum(control) / len(control)
                estimates.append(est)

        if len(estimates) < 2:
            return 0.0
        return float(_variance(estimates) ** 0.5)


# ── Sensitivity Analysis ──────────────────────────────────────


class SensitivityAnalyzer:
    """Rosenbaum-style sensitivity analysis for unmeasured confounding.

    Tests how strong an unmeasured confounder would need to be
    (gamma) to explain away the estimated effect.
    """

    def analyze(
        self,
        data: list[dict[str, float]],
        treatment: str,
        outcome: str,
        gamma_values: list[float] | None = None,
    ) -> list[SensitivityResult]:
        """Run sensitivity analysis across gamma values."""
        if gamma_values is None:
            gamma_values = [1.0, 1.5, 2.0, 3.0]

        treated = [d[outcome] for d in data if d[treatment] >= 0.5]
        control = [d[outcome] for d in data if d[treatment] < 0.5]

        if not treated or not control:
            return [
                SensitivityResult(
                    gamma=g,
                    p_value_upper=1.0,
                    p_value_lower=1.0,
                    robust=False,
                )
                for g in gamma_values
            ]

        mean_t = sum(treated) / len(treated)
        mean_c = sum(control) / len(control)
        diff = mean_t - mean_c

        n_t = len(treated)
        n_c = len(control)
        pooled_var = (_variance(treated) + _variance(control)) / 2
        se = (pooled_var * (1 / n_t + 1 / n_c)) ** 0.5 if pooled_var > 0 else 1e-10

        results: list[SensitivityResult] = []
        for gamma in gamma_values:
            # Rosenbaum bounds: shift the effect by log(gamma)
            shift = math.log(max(gamma, 1.001))
            z_upper = (diff - shift) / max(se, 1e-10)
            z_lower = (diff + shift) / max(se, 1e-10)

            p_upper = self._normal_cdf(-abs(z_upper))
            p_lower = self._normal_cdf(-abs(z_lower))

            results.append(
                SensitivityResult(
                    gamma=gamma,
                    p_value_upper=p_upper,
                    p_value_lower=p_lower,
                    robust=(p_upper < 0.05),
                )
            )

        return results

    @staticmethod
    def _normal_cdf(z: float) -> float:
        """Approximate standard normal CDF using Abramowitz & Stegun."""
        if z < -8.0:
            return 0.0
        if z > 8.0:
            return 1.0
        a1, a2, a3, a4, a5 = (
            0.254829592,
            -0.284496736,
            1.421413741,
            -1.453152027,
            1.061405429,
        )
        p = 0.3275911
        sign = 1 if z >= 0 else -1
        z_abs = abs(z)
        t = 1.0 / (1.0 + p * z_abs)
        y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * math.exp(
            -z_abs * z_abs / 2
        )
        return 0.5 * (1.0 + sign * y)


# ── Mediation Analysis ────────────────────────────────────────


class MediationAnalyzer:
    """Causal mediation analysis for direct/indirect effects.

    Decomposes total effect of X on Y through mediator M into:
    - Natural Direct Effect (NDE): X → Y not through M
    - Natural Indirect Effect (NIE): X → M → Y
    """

    def __init__(self, dag: WorldDAG) -> None:
        self._dag = dag

    def analyze(
        self,
        data: list[dict[str, float]],
        treatment: str,
        outcome: str,
        mediator: str,
    ) -> MediationResult:
        """Estimate mediation decomposition via regression."""
        x = [d[treatment] for d in data]
        m = [d[mediator] for d in data]
        y = [d[outcome] for d in data]

        # Total effect: regress Y on X
        total = _regression_coefficient(x, y)

        # Direct effect: regress Y on X controlling for M
        direct = _partial_regression_coefficient(x, y, [m])

        indirect = total - direct

        prop = 0.0
        if abs(total) > 1e-10:
            prop = abs(indirect / total)
            prop = min(prop, 1.0)

        return MediationResult(
            total_effect=total,
            direct_effect=direct,
            indirect_effect=indirect,
            proportion_mediated=prop,
            mediator=mediator,
        )
