"""Tests for the counterfactual reasoning engine."""

from __future__ import annotations

import random

import pytest

from orxaq.intelligence.counterfactual import (
    CounterfactualEngine,
    CounterfactualQuery,
    CounterfactualResult,
    NecessaryCausation,
    SufficientCausation,
)
from orxaq.kernel import CausalEdge, CausalRole, Entity, VariableType, WorldDAG


def _entity(eid: str) -> Entity:
    return Entity(
        entity_id=eid,
        name=eid,
        causal_role=CausalRole.EXOGENOUS,
        variable_type=VariableType.STATIC,
    )


def _edge(src: str, tgt: str) -> CausalEdge:
    return CausalEdge(source=src, target=tgt, confidence=0.9)


@pytest.fixture()
def simple_dag() -> WorldDAG:
    """X → Y (direct effect)."""
    dag = WorldDAG()
    for n in ["X", "Y"]:
        dag.add_entity(_entity(n))
    dag.add_edge(_edge("X", "Y"))
    return dag


@pytest.fixture()
def chain_dag() -> WorldDAG:
    """X → M → Y."""
    dag = WorldDAG()
    for n in ["X", "M", "Y"]:
        dag.add_entity(_entity(n))
    dag.add_edge(_edge("X", "M"))
    dag.add_edge(_edge("M", "Y"))
    return dag


@pytest.fixture()
def confounded_dag() -> WorldDAG:
    """Z → X, Z → Y, X → Y."""
    dag = WorldDAG()
    for n in ["X", "Y", "Z"]:
        dag.add_entity(_entity(n))
    dag.add_edge(_edge("Z", "X"))
    dag.add_edge(_edge("Z", "Y"))
    dag.add_edge(_edge("X", "Y"))
    return dag


def _gen_binary_treatment(n: int = 300, seed: int = 42) -> list[dict[str, float]]:
    """Generate data with binary X and Y = 0.6*X + noise."""
    rng = random.Random(seed)
    data = []
    for _ in range(n):
        x = rng.choice([0.0, 1.0])
        y = float(0.6 * x + rng.gauss(0, 0.3) > 0.3)
        data.append({"X": x, "Y": y})
    return data


def _gen_chain(n: int = 300, seed: int = 42) -> list[dict[str, float]]:
    rng = random.Random(seed)
    data = []
    for _ in range(n):
        x = rng.choice([0.0, 1.0])
        m = 0.8 * x + rng.gauss(0, 0.2)
        y = float(0.7 * m + rng.gauss(0, 0.2) > 0.3)
        data.append({"X": x, "M": m, "Y": y})
    return data


def _gen_confounded(n: int = 300, seed: int = 42) -> list[dict[str, float]]:
    rng = random.Random(seed)
    data = []
    for _ in range(n):
        z = rng.gauss(0, 1)
        x = float(0.5 * z + rng.gauss(0, 0.3) > 0)
        y = float(0.3 * x + 0.4 * z + rng.gauss(0, 0.3) > 0.3)
        data.append({"X": x, "Y": y, "Z": z})
    return data


class TestCounterfactualQuery:
    def test_construction(self) -> None:
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={"X": 1.0, "Z": 0.5},
            intervention_value=0.0,
        )
        assert q.treatment == "X"
        assert q.outcome == "Y"
        assert q.factual_value is None

    def test_with_factual_value(self) -> None:
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={"X": 1.0},
            intervention_value=0.0,
            factual_value=0.8,
        )
        assert q.factual_value == 0.8

    def test_frozen(self) -> None:
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={},
            intervention_value=1.0,
        )
        with pytest.raises(AttributeError):
            q.treatment = "Z"  # type: ignore[misc]


class TestCounterfactualResult:
    def test_fields(self) -> None:
        r = CounterfactualResult(
            factual_outcome=0.0,
            counterfactual_outcome=0.6,
            individual_treatment_effect=0.6,
            probability_of_necessity=0.8,
            probability_of_sufficiency=0.7,
        )
        assert r.individual_treatment_effect == 0.6
        assert r.probability_of_necessity == 0.8


class TestCounterfactualEngine:
    def test_simple_query(self, simple_dag: WorldDAG) -> None:
        engine = CounterfactualEngine(simple_dag)
        data = _gen_binary_treatment()
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={"X": 1.0, "Y": 1.0},
            intervention_value=0.0,
        )
        result = engine.query(data, q)
        assert isinstance(result, CounterfactualResult)
        # ITE should reflect removing treatment
        assert result.individual_treatment_effect != 0.0

    def test_positive_ite(self, simple_dag: WorldDAG) -> None:
        engine = CounterfactualEngine(simple_dag)
        data = _gen_binary_treatment()
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={"X": 0.0, "Y": 0.0},
            intervention_value=1.0,
        )
        result = engine.query(data, q)
        # Adding treatment should increase Y
        assert result.individual_treatment_effect > 0

    def test_factual_override(self, simple_dag: WorldDAG) -> None:
        engine = CounterfactualEngine(simple_dag)
        data = _gen_binary_treatment()
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={"X": 1.0},
            intervention_value=0.0,
            factual_value=0.9,
        )
        result = engine.query(data, q)
        assert result.factual_outcome == 0.9

    def test_chain_dag_query(self, chain_dag: WorldDAG) -> None:
        engine = CounterfactualEngine(chain_dag)
        data = _gen_chain()
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={"X": 1.0, "M": 0.8, "Y": 1.0},
            intervention_value=0.0,
        )
        result = engine.query(data, q)
        assert isinstance(result, CounterfactualResult)

    def test_confounded_query(self, confounded_dag: WorldDAG) -> None:
        engine = CounterfactualEngine(confounded_dag)
        data = _gen_confounded()
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={"X": 1.0, "Y": 1.0, "Z": 0.5},
            intervention_value=0.0,
        )
        result = engine.query(data, q)
        assert isinstance(result, CounterfactualResult)

    def test_pn_bounded(self, simple_dag: WorldDAG) -> None:
        engine = CounterfactualEngine(simple_dag)
        data = _gen_binary_treatment()
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={"X": 1.0, "Y": 1.0},
            intervention_value=0.0,
        )
        result = engine.query(data, q)
        assert 0.0 <= result.probability_of_necessity <= 1.0

    def test_ps_bounded(self, simple_dag: WorldDAG) -> None:
        engine = CounterfactualEngine(simple_dag)
        data = _gen_binary_treatment()
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={"X": 0.0, "Y": 0.0},
            intervention_value=1.0,
        )
        result = engine.query(data, q)
        assert 0.0 <= result.probability_of_sufficiency <= 1.0

    def test_no_parents_outcome(self) -> None:
        """Y has no parents → counterfactual uses prediction."""
        dag = WorldDAG()
        dag.add_entity(_entity("X"))
        dag.add_entity(_entity("Y"))
        # No edges
        engine = CounterfactualEngine(dag)
        data = [{"X": 1.0, "Y": 0.5}] * 10
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={"X": 1.0, "Y": 0.5},
            intervention_value=0.0,
        )
        result = engine.query(data, q)
        assert isinstance(result, CounterfactualResult)


class TestNecessaryCausation:
    def test_but_for_bounded(self) -> None:
        nc = NecessaryCausation()
        data = _gen_binary_treatment()
        result = nc.but_for_test(data, "X", "Y")
        assert 0.0 <= result <= 1.0

    def test_but_for_strong_effect(self) -> None:
        """Strong treatment effect → high PN."""
        rng = random.Random(42)
        data = []
        for _ in range(200):
            x = rng.choice([0.0, 1.0])
            y = float(x > 0.5)  # Y = X deterministically
            data.append({"X": x, "Y": y})
        nc = NecessaryCausation()
        result = nc.but_for_test(data, "X", "Y")
        assert result > 0.5

    def test_but_for_no_treated(self) -> None:
        data = [{"X": 0.0, "Y": 0.0}] * 10
        nc = NecessaryCausation()
        result = nc.but_for_test(data, "X", "Y")
        assert result == 0.0

    def test_but_for_all_treated_negative(self) -> None:
        """All treated but outcome is 0 → no treated-positive → PN = 0."""
        data = [{"X": 1.0, "Y": 0.0}] * 10
        nc = NecessaryCausation()
        result = nc.but_for_test(data, "X", "Y")
        assert result == 0.0


class TestSufficientCausation:
    def test_would_cause_bounded(self) -> None:
        sc = SufficientCausation()
        data = _gen_binary_treatment()
        result = sc.would_cause_test(data, "X", "Y")
        assert 0.0 <= result <= 1.0

    def test_would_cause_strong_effect(self) -> None:
        """Strong treatment effect → high PS."""
        rng = random.Random(42)
        data = []
        for _ in range(200):
            x = rng.choice([0.0, 1.0])
            y = float(x > 0.5)
            data.append({"X": x, "Y": y})
        sc = SufficientCausation()
        result = sc.would_cause_test(data, "X", "Y")
        assert result > 0.5

    def test_would_cause_no_untreated(self) -> None:
        data = [{"X": 1.0, "Y": 1.0}] * 10
        sc = SufficientCausation()
        result = sc.would_cause_test(data, "X", "Y")
        assert result == 0.0

    def test_would_cause_all_untreated_positive(self) -> None:
        """All untreated with positive outcome → no untreated-negative → PS = 0."""
        data = [{"X": 0.0, "Y": 1.0}] * 10
        sc = SufficientCausation()
        result = sc.would_cause_test(data, "X", "Y")
        assert result == 0.0


# ── Targeted coverage (lines 98, 135, 188) ──────────────────


class TestCounterfactualEdgeCases:
    """Cover lines 98, 135, 188: empty data, small data, no treated_positive."""

    def test_predict_empty_data(self) -> None:
        """Cover line 98: _predict with empty data returns 0.0."""
        dag = WorldDAG()
        dag.add_entity(_entity("X"))
        dag.add_entity(_entity("Y"))
        dag.add_edge(CausalEdge(source="X", target="Y"))
        engine = CounterfactualEngine(dag)
        result = engine._predict([], {"X": 1.0}, "Y")
        assert result == 0.0

    def test_query_small_data_fallback(self) -> None:
        """Cover line 135: n<2 falls back to prediction."""
        dag = WorldDAG()
        dag.add_entity(_entity("X"))
        dag.add_entity(_entity("Y"))
        dag.add_edge(CausalEdge(source="X", target="Y"))
        engine = CounterfactualEngine(dag)
        data = [{"X": 1.0, "Y": 0.8}]
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={"X": 1.0, "Y": 0.8},
            intervention_value=0.0,
        )
        result = engine.query(data, q)
        assert isinstance(result, CounterfactualResult)

    def test_but_for_no_treated_positive(self) -> None:
        """Cover line 260: no treated with positive outcome in but_for_test."""
        nc = NecessaryCausation()
        # All treated have negative outcome
        data = [{"X": 1.0, "Y": 0.0}, {"X": 1.0, "Y": 0.0}]
        result = nc.but_for_test(data, "X", "Y")
        assert result == 0.0

    def test_query_no_treated_positive(self) -> None:
        """Cover line 188: _probability_of_necessity with no treated-positive."""
        dag = WorldDAG()
        dag.add_entity(_entity("X"))
        dag.add_entity(_entity("Y"))
        dag.add_edge(CausalEdge(source="X", target="Y"))
        engine = CounterfactualEngine(dag)
        # No unit has both X >= 0.5 AND Y >= 0.5
        data = [
            {"X": 0.0, "Y": 0.0},
            {"X": 0.0, "Y": 0.0},
            {"X": 1.0, "Y": 0.0},
        ]
        q = CounterfactualQuery(
            treatment="X",
            outcome="Y",
            evidence={"X": 1.0, "Y": 0.0},
            intervention_value=0.0,
        )
        result = engine.query(data, q)
        assert result.probability_of_necessity == 0.0
