"""Counterfactual reasoning engine for causal what-if queries.

Implements Pearl's three-step counterfactual algorithm:
1. Abduction: infer exogenous noise from observed evidence
2. Action: intervene on the treatment variable
3. Prediction: propagate effects through the structural model

Also computes Probability of Necessity (PN) and Probability of
Sufficiency (PS) for legal/regulatory causation analysis.

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

from __future__ import annotations

from dataclasses import dataclass

from orxaq.intelligence.intervention import (
    _regression_coefficient,
)
from orxaq.kernel import WorldDAG


@dataclass(frozen=True)
class CounterfactualQuery:
    """A counterfactual query: what would Y be if X had been x'?"""

    treatment: str
    outcome: str
    evidence: dict[str, float]
    intervention_value: float
    factual_value: float | None = None


@dataclass(frozen=True)
class CounterfactualResult:
    """Result of a counterfactual query."""

    factual_outcome: float
    counterfactual_outcome: float
    individual_treatment_effect: float
    probability_of_necessity: float
    probability_of_sufficiency: float


class CounterfactualEngine:
    """Three-step counterfactual reasoning using linear structural models.

    Uses regression-based structural equation estimation:
    1. Fit Y = alpha + beta*X + gamma*Z + epsilon from data
    2. For a specific unit, infer epsilon from observed values
    3. Compute counterfactual Y' under do(X=x') using same epsilon
    """

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

    def query(
        self,
        data: list[dict[str, float]],
        q: CounterfactualQuery,
    ) -> CounterfactualResult:
        """Answer a counterfactual query using the three-step method."""
        # Step 1: Determine factual outcome
        factual = q.factual_value
        if factual is None:
            factual = self._predict(data, q.evidence, q.outcome)

        # Step 2: Compute counterfactual outcome under intervention
        cf_evidence = dict(q.evidence)
        cf_evidence[q.treatment] = q.intervention_value
        cf_outcome = self._structural_counterfactual(
            data, q.evidence, cf_evidence, q.outcome, q.treatment
        )

        ite = cf_outcome - factual

        # Step 3: PN and PS via regression-based bounds
        pn = self._probability_of_necessity(data, q.treatment, q.outcome)
        ps = self._probability_of_sufficiency(data, q.treatment, q.outcome)

        return CounterfactualResult(
            factual_outcome=factual,
            counterfactual_outcome=cf_outcome,
            individual_treatment_effect=ite,
            probability_of_necessity=pn,
            probability_of_sufficiency=ps,
        )

    def _predict(
        self,
        data: list[dict[str, float]],
        evidence: dict[str, float],
        outcome: str,
    ) -> float:
        """Predict outcome using nearest-neighbor matching on evidence."""
        if not data:
            return 0.0

        best_dist = float("inf")
        best_y = 0.0

        for row in data:
            dist = sum((row.get(k, 0.0) - v) ** 2 for k, v in evidence.items())
            if dist < best_dist:
                best_dist = dist
                best_y = row.get(outcome, 0.0)

        return best_y

    def _structural_counterfactual(
        self,
        data: list[dict[str, float]],
        factual_evidence: dict[str, float],
        cf_evidence: dict[str, float],
        outcome: str,
        treatment: str,
    ) -> float:
        """Three-step structural counterfactual.

        1. Estimate structural coefficients from data
        2. Abduction: infer residual for this unit
        3. Prediction: compute Y' under intervention
        """
        parents = self._dag.parents(outcome)
        if not parents:
            return self._predict(data, cf_evidence, outcome)

        parent_list = sorted(parents)

        # Estimate structural coefficients: Y = intercept + sum(beta_i * X_i)
        y_vals = [d.get(outcome, 0.0) for d in data]
        n = len(data)
        if n < 2:
            return self._predict(data, cf_evidence, outcome)

        y_mean = sum(y_vals) / n
        intercept = y_mean
        betas: dict[str, float] = {}

        for p in parent_list:
            x_vals = [d.get(p, 0.0) for d in data]
            beta = _regression_coefficient(x_vals, y_vals)
            betas[p] = beta
            x_mean = sum(x_vals) / n
            intercept -= beta * x_mean

        # Abduction: compute residual for factual unit
        factual_predicted = intercept + sum(
            betas.get(p, 0.0) * factual_evidence.get(p, 0.0) for p in parent_list
        )
        factual_y = factual_evidence.get(outcome, self._predict(data, factual_evidence, outcome))
        residual = factual_y - factual_predicted

        # Prediction: counterfactual outcome using intervened values
        cf_y = (
            intercept
            + sum(betas.get(p, 0.0) * cf_evidence.get(p, 0.0) for p in parent_list)
            + residual
        )

        return cf_y

    def _probability_of_necessity(
        self,
        data: list[dict[str, float]],
        treatment: str,
        outcome: str,
    ) -> float:
        """P(Y_{x=0} = 0 | X = 1, Y = 1).

        Among treated with positive outcome: what fraction would have
        had negative outcome without treatment?

        Estimated via regression: for each treated unit with Y=1,
        predict Y under do(X=0). If predicted Y < 0.5, count as
        "would have been negative".
        """
        x_vals = [d.get(treatment, 0.0) for d in data]
        y_vals = [d.get(outcome, 0.0) for d in data]
        beta = _regression_coefficient(x_vals, y_vals)

        treated_positive = [
            d for d in data if d.get(treatment, 0.0) >= 0.5 and d.get(outcome, 0.0) >= 0.5
        ]

        if not treated_positive:
            return 0.0

        necessary_count = 0
        for unit in treated_positive:
            # Counterfactual: what Y would be if X were 0
            cf_y = unit.get(outcome, 0.0) - beta * unit.get(treatment, 0.0)
            if cf_y < 0.5:
                necessary_count += 1

        return necessary_count / len(treated_positive)

    def _probability_of_sufficiency(
        self,
        data: list[dict[str, float]],
        treatment: str,
        outcome: str,
    ) -> float:
        """P(Y_{x=1} = 1 | X = 0, Y = 0).

        Among untreated with negative outcome: what fraction would have
        had positive outcome with treatment?

        Estimated via regression: for each untreated unit with Y=0,
        predict Y under do(X=1). If predicted Y >= 0.5, count as
        "would have been positive".
        """
        x_vals = [d.get(treatment, 0.0) for d in data]
        y_vals = [d.get(outcome, 0.0) for d in data]
        beta = _regression_coefficient(x_vals, y_vals)

        untreated_negative = [
            d for d in data if d.get(treatment, 0.0) < 0.5 and d.get(outcome, 0.0) < 0.5
        ]

        if not untreated_negative:
            return 0.0

        sufficient_count = 0
        for unit in untreated_negative:
            # Counterfactual: what Y would be if X were 1
            cf_y = unit.get(outcome, 0.0) + beta * (1.0 - unit.get(treatment, 0.0))
            if cf_y >= 0.5:
                sufficient_count += 1

        return sufficient_count / len(untreated_negative)


class NecessaryCausation:
    """But-for test: was the treatment necessary for the outcome?

    Uses regression to estimate counterfactual outcomes.
    """

    def but_for_test(
        self,
        data: list[dict[str, float]],
        treatment: str,
        outcome: str,
    ) -> float:
        """Fraction of treated-positive units that needed the treatment.

        Returns P(Y_0 < 0.5 | X >= 0.5, Y >= 0.5), estimated by
        computing the counterfactual Y under do(X=0) for each unit.
        """
        x_vals = [d.get(treatment, 0.0) for d in data]
        y_vals = [d.get(outcome, 0.0) for d in data]
        beta = _regression_coefficient(x_vals, y_vals)

        treated_positive = [
            d for d in data if d.get(treatment, 0.0) >= 0.5 and d.get(outcome, 0.0) >= 0.5
        ]

        if not treated_positive:
            return 0.0

        count = 0
        for unit in treated_positive:
            cf_y = unit[outcome] - beta * unit[treatment]
            if cf_y < 0.5:
                count += 1

        return count / len(treated_positive)


class SufficientCausation:
    """Would-cause test: was the treatment sufficient for the outcome?

    Uses regression to estimate counterfactual outcomes.
    """

    def would_cause_test(
        self,
        data: list[dict[str, float]],
        treatment: str,
        outcome: str,
    ) -> float:
        """Fraction of untreated-negative units that would benefit.

        Returns P(Y_1 >= 0.5 | X < 0.5, Y < 0.5), estimated by
        computing the counterfactual Y under do(X=1) for each unit.
        """
        x_vals = [d.get(treatment, 0.0) for d in data]
        y_vals = [d.get(outcome, 0.0) for d in data]
        beta = _regression_coefficient(x_vals, y_vals)

        untreated_negative = [
            d for d in data if d.get(treatment, 0.0) < 0.5 and d.get(outcome, 0.0) < 0.5
        ]

        if not untreated_negative:
            return 0.0

        count = 0
        for unit in untreated_negative:
            cf_y = unit[outcome] + beta * (1.0 - unit[treatment])
            if cf_y >= 0.5:
                count += 1

        return count / len(untreated_negative)
