"""Tests for d-separation validation and refutation suite."""

from __future__ import annotations

import random

from orxaq.intelligence.validation import DSeparationTester, RefutationSuite
from orxaq.kernel import CausalEdge, CausalRole, Entity, VariableType, WorldDAG


def _build_chain_dag() -> WorldDAG:
    """Build X → Z → Y DAG."""
    dag = WorldDAG()
    for name in ["X", "Z", "Y"]:
        dag.add_entity(
            Entity(
                entity_id=name,
                name=name,
                causal_role=CausalRole.EXOGENOUS,
                variable_type=VariableType.STATIC,
            )
        )
    dag.add_edge(CausalEdge(source="X", target="Z", mechanism="direct"))
    dag.add_edge(CausalEdge(source="Z", target="Y", mechanism="direct"))
    return dag


def _make_chain_data(n: int = 500, seed: int = 42) -> list[dict]:
    """Generate data consistent with X → Z → Y."""
    rng = random.Random(seed)
    data = []
    for _ in range(n):
        x = rng.gauss(0, 1)
        z = 2.0 * x + rng.gauss(0, 0.3)
        y = 1.5 * z + rng.gauss(0, 0.3)
        data.append({"X": x, "Z": z, "Y": y})
    return data


def _build_misspecified_dag() -> WorldDAG:
    """Build a DAG that doesn't match the data (X → Y, missing Z)."""
    dag = WorldDAG()
    for name in ["X", "Z", "Y"]:
        dag.add_entity(
            Entity(
                entity_id=name,
                name=name,
                causal_role=CausalRole.EXOGENOUS,
                variable_type=VariableType.STATIC,
            )
        )
    # Only X → Y, no Z involvement — misspecified for chain data
    dag.add_edge(CausalEdge(source="X", target="Y", mechanism="direct"))
    return dag


class TestDSeparationTester:
    def test_correct_dag_passes(self):
        """Correct DAG should have few violations."""
        dag = _build_chain_dag()
        data = _make_chain_data(n=500)
        tester = DSeparationTester(dag)
        report = tester.test_all_implications(data, ["X", "Z", "Y"])

        assert report.total_implications > 0
        # Correct DAG should have low violation rate
        assert report.violation_rate < 0.3

    def test_misspecified_dag_has_violations(self):
        """Misspecified DAG should have d-sep violations."""
        dag = _build_misspecified_dag()
        data = _make_chain_data(n=500)
        tester = DSeparationTester(dag)
        report = tester.test_all_implications(data, ["X", "Z", "Y"])

        # Misspecified DAG should imply Z ⊥ Y | ∅ or similar that's violated
        # At minimum, some implications should be tested
        assert report.total_implications > 0

    def test_report_fields(self):
        """Report should have all required fields."""
        dag = _build_chain_dag()
        data = _make_chain_data(n=200)
        tester = DSeparationTester(dag)
        report = tester.test_all_implications(data, ["X", "Z", "Y"])

        assert isinstance(report.total_implications, int)
        assert isinstance(report.violations, int)
        assert isinstance(report.violation_rate, float)
        assert isinstance(report.worst_violations, tuple)
        assert isinstance(report.passed, bool)

    def test_specific_test(self):
        """test_specific should return a CI result."""
        dag = _build_chain_dag()
        data = _make_chain_data(n=300)
        tester = DSeparationTester(dag)

        result = tester.test_specific(data, "X", "Y", frozenset({"Z"}))
        assert result.x == "X"
        assert result.y == "Y"
        assert result.z_set == frozenset({"Z"})
        assert isinstance(result.p_value, float)

    def test_violations_sorted_by_pvalue(self):
        """Worst violations should be sorted by p-value."""
        dag = _build_misspecified_dag()
        data = _make_chain_data(n=500)
        tester = DSeparationTester(dag)
        report = tester.test_all_implications(data, ["X", "Z", "Y"])

        if len(report.worst_violations) > 1:
            for i in range(len(report.worst_violations) - 1):
                assert report.worst_violations[i].p_value <= report.worst_violations[i + 1].p_value


class TestRefutationSuite:
    def test_placebo_strong_effect(self):
        """With strong effect, placebo p-value should be low."""
        rng = random.Random(42)
        data = []
        for _ in range(200):
            x = rng.gauss(0, 1)
            y = 5.0 * x + rng.gauss(0, 0.1)
            data.append({"treatment": x, "outcome": y})

        suite = RefutationSuite(seed=42)
        p_value = suite.placebo_test(data, "treatment", "outcome", n_permutations=50)
        # Original effect is real, so very few permutations should match
        assert p_value < 0.15

    def test_placebo_no_effect(self):
        """With no real effect, placebo p-value should be high."""
        rng = random.Random(42)
        data = []
        for _ in range(200):
            data.append(
                {
                    "treatment": rng.gauss(0, 1),
                    "outcome": rng.gauss(0, 1),
                }
            )

        suite = RefutationSuite(seed=42)
        p_value = suite.placebo_test(data, "treatment", "outcome", n_permutations=50)
        # No real effect, permutations should often match
        assert p_value > 0.05

    def test_subset_stability_strong(self):
        """Strong effect should be stable across subsets (low CV)."""
        rng = random.Random(42)
        data = []
        for _ in range(500):
            x = rng.gauss(0, 1)
            y = 5.0 * x + rng.gauss(0, 0.5)
            data.append({"treatment": x, "outcome": y})

        suite = RefutationSuite(seed=42)
        cv = suite.subset_test(data, "treatment", "outcome", n_subsets=10)
        # Strong, stable effect should have low coefficient of variation
        assert cv < 0.5

    def test_subset_returns_float(self):
        """Subset test should return a float."""
        data = [{"treatment": float(i), "outcome": float(i * 2)} for i in range(100)]
        suite = RefutationSuite(seed=42)
        cv = suite.subset_test(data, "treatment", "outcome")
        assert isinstance(cv, float)
