"""D-separation testing and causal refutation suite.

Tests whether conditional independence implications of a DAG hold in data.
Violations indicate model misspecification. Complemented by placebo and
subset refutation tests. Stdlib-only (ADR-007, ADR-019).
"""

from __future__ import annotations

import random
from dataclasses import dataclass
from itertools import combinations

from orxaq.intelligence.independence import CIResult, PartialCorrelationTest
from orxaq.kernel import WorldDAG


@dataclass(frozen=True)
class DSepImplication:
    """A d-separation implication and its test result."""

    x: str
    y: str
    z_set: frozenset[str]
    dag_implies_independent: bool
    data_independent: bool
    violated: bool
    p_value: float


@dataclass(frozen=True)
class DSepReport:
    """Report from testing all d-separation implications."""

    total_implications: int
    violations: int
    violation_rate: float
    worst_violations: tuple[DSepImplication, ...]
    passed: bool


class DSeparationTester:
    """Test d-separation implications of a DAG against data.

    For each pair (X, Y) and conditioning set Z where the DAG implies
    X is independent of Y given Z, test this against data. Violations
    are evidence of model misspecification.
    """

    def __init__(
        self,
        dag: WorldDAG,
        ci_test: PartialCorrelationTest | None = None,
    ) -> None:
        self._dag = dag
        self._ci_test = ci_test or PartialCorrelationTest()

    def test_all_implications(
        self,
        data: list[dict],
        variables: list[str] | None = None,
        alpha: float = 0.05,
        max_cond_size: int = 2,
        max_tests: int = 200,
    ) -> DSepReport:
        """Test all d-separation implications up to max_cond_size."""
        if variables is None:
            variables = sorted(self._dag._entities.keys())

        implications: list[DSepImplication] = []
        test_count = 0

        for i, x in enumerate(variables):
            for y in variables[i + 1 :]:
                if test_count >= max_tests:
                    break

                # Test with empty conditioning set
                impl = self._test_implication(data, x, y, frozenset(), alpha)
                implications.append(impl)
                test_count += 1

                # Test with conditioning sets of increasing size
                other_vars = [v for v in variables if v != x and v != y]
                for size in range(1, min(max_cond_size + 1, len(other_vars) + 1)):
                    if test_count >= max_tests:
                        break
                    for z_tuple in combinations(other_vars, size):
                        if test_count >= max_tests:
                            break
                        z_set = frozenset(z_tuple)
                        impl = self._test_implication(data, x, y, z_set, alpha)
                        implications.append(impl)
                        test_count += 1

        violations = [imp for imp in implications if imp.violated]
        violation_count = len(violations)
        total = len(implications)
        rate = violation_count / total if total > 0 else 0.0

        # Sort violations by p-value (most significant first)
        worst = sorted(violations, key=lambda v: v.p_value)[:10]

        return DSepReport(
            total_implications=total,
            violations=violation_count,
            violation_rate=rate,
            worst_violations=tuple(worst),
            passed=rate < 0.10,  # <10% violation rate
        )

    def _test_implication(
        self,
        data: list[dict],
        x: str,
        y: str,
        z_set: frozenset[str],
        alpha: float,
    ) -> DSepImplication:
        """Test a single d-separation implication."""
        dag_independent = self._dag.d_separated(x, y, z_set)
        ci_result = self._ci_test.test(data, x, y, z_set, alpha)

        # Violation: DAG says independent but data says dependent
        violated = dag_independent and not ci_result.independent

        return DSepImplication(
            x=x,
            y=y,
            z_set=z_set,
            dag_implies_independent=dag_independent,
            data_independent=ci_result.independent,
            violated=violated,
            p_value=ci_result.p_value,
        )

    def test_specific(
        self,
        data: list[dict],
        x: str,
        y: str,
        z_set: frozenset[str],
        alpha: float = 0.05,
    ) -> CIResult:
        """Test a specific conditional independence claim."""
        return self._ci_test.test(data, x, y, z_set, alpha)


class RefutationSuite:
    """Standard causal inference refutation tests."""

    def __init__(self, seed: int = 42) -> None:
        self._rng = random.Random(seed)

    def placebo_test(
        self,
        data: list[dict],
        treatment: str,
        outcome: str,
        n_permutations: int = 100,
    ) -> float:
        """Placebo test: replace treatment with random noise.

        The causal effect should approach 0 when treatment is random.
        Returns the fraction of permuted effects that are smaller than
        the original effect — should be close to 0 (most permuted
        effects should be smaller → original is real).
        """
        ci_test = PartialCorrelationTest()

        # Original effect
        original = ci_test.test(data, treatment, outcome)
        original_stat = abs(original.statistic)

        # Permutation distribution
        count_larger = 0
        treatment_vals = [row.get(treatment, 0) for row in data]

        for _ in range(n_permutations):
            shuffled = treatment_vals[:]
            self._rng.shuffle(shuffled)
            permuted_data = []
            for i, row in enumerate(data):
                new_row = dict(row)
                new_row[treatment] = shuffled[i]
                permuted_data.append(new_row)

            perm_result = ci_test.test(permuted_data, treatment, outcome)
            if abs(perm_result.statistic) >= original_stat:
                count_larger += 1

        # p-value: fraction of permuted effects >= original
        return count_larger / n_permutations

    def subset_test(
        self,
        data: list[dict],
        treatment: str,
        outcome: str,
        n_subsets: int = 5,
    ) -> float:
        """Subset test: effect should be stable across random subsets.

        Returns the coefficient of variation of the effect across subsets.
        Low CV (< 0.5) suggests the effect is robust.
        """
        ci_test = PartialCorrelationTest()
        n = len(data)
        subset_size = max(n // 2, 10)

        effects: list[float] = []
        for _ in range(n_subsets):
            indices = self._rng.sample(range(n), min(subset_size, n))
            subset = [data[i] for i in indices]
            result = ci_test.test(subset, treatment, outcome)
            effects.append(abs(result.statistic))

        if not effects:
            return 0.0

        mean_effect = sum(effects) / len(effects)
        if mean_effect < 1e-10:
            return 0.0

        variance = sum((e - mean_effect) ** 2 for e in effects) / len(effects)
        std = variance**0.5

        return float(std / mean_effect)  # coefficient of variation
