"""Tests for inference/confounders.py.

Hand-built conceptual models with known DAG structures.
"""

from rtisanepy import Unit, continuous, causes, ConceptualModel
from rtisanepy.variables import UnobservedVariable
from rtisanepy.relationships import Causes
from rtisanepy.conceptual_model import Assumption, Hypothesis
from rtisanepy.inference.confounders import infer_confounders


def _quick_cm(*assumed_causes, hypothesis):
    """Build a ConceptualModel from (cause, effect) pairs + one hypothesis."""
    cm = ConceptualModel()
    for c, e in assumed_causes:
        cm.assume(causes(c, e))
    cm.hypothesize(causes(hypothesis[0], hypothesis[1]))
    return cm


class TestModel1CommonParent:
    def test_common_observed_parent(self):
        u = Unit("s", cardinality=10)
        x = continuous(u, "X")
        y = continuous(u, "Y")
        z = continuous(u, "Z")  # common cause
        cm = _quick_cm((z, x), (z, y), hypothesis=(x, y))
        conf = infer_confounders(cm, x, y)
        assert "Z" in conf

    def test_no_common_parent(self):
        u = Unit("s", cardinality=10)
        x = continuous(u, "X")
        y = continuous(u, "Y")
        cm = ConceptualModel().hypothesize(causes(x, y))
        conf = infer_confounders(cm, x, y)
        assert conf == []


class TestModel2UnobservedAncestorMediated:
    def test_unobserved_mediates_via_iv_parent(self):
        """U (unobserved) -> Z -> X, U -> Y  =>  Z is confounder."""
        u = Unit("s", cardinality=10)
        x = continuous(u, "X")
        y = continuous(u, "Y")
        z = continuous(u, "Z")
        unobs = UnobservedVariable("U")

        cm = ConceptualModel()
        # U -> Z
        cm.variables["U"] = unobs
        cm.assume(causes(unobs, z))
        # Z -> X
        cm.assume(causes(z, x))
        # U -> Y
        cm.assume(causes(unobs, y))
        # Hypothesis: X -> Y
        cm.hypothesize(causes(x, y))

        conf = infer_confounders(cm, x, y)
        assert "Z" in conf


class TestModel8NeutralControl:
    def test_parent_of_y_unrelated_to_x(self):
        """W -> Y only, X -> Y  =>  W is neutral control."""
        u = Unit("s", cardinality=10)
        x = continuous(u, "X")
        y = continuous(u, "Y")
        w = continuous(u, "W")

        cm = (
            ConceptualModel()
            .assume(causes(w, y))
            .hypothesize(causes(x, y))
        )
        conf = infer_confounders(cm, x, y)
        assert "W" in conf


class TestMediators:
    def test_mediator_not_added_as_confounder(self):
        """X -> M -> Y: M is a mediator, not a confounder on its own."""
        u = Unit("s", cardinality=10)
        x = continuous(u, "X")
        y = continuous(u, "Y")
        m = continuous(u, "M")

        cm = (
            ConceptualModel()
            .assume(causes(x, m))
            .assume(causes(m, y))
            .hypothesize(causes(x, y))
        )
        conf = infer_confounders(cm, x, y)
        # M should NOT appear as a confounder
        assert "M" not in conf
