"""End-to-end tests: build conceptual model → query → StatisticalModel."""

from rtisanepy import (
    Unit, continuous, categories, counts,
    causes, relates, equals, increases,
    ConceptualModel, StatisticalModel,
)


class TestBasicQuery:
    def test_simple_cause(self):
        student = Unit("student", cardinality=100)
        tutoring = categories(student, "Tutoring", cardinality=2)
        score = continuous(student, "Score")

        cm = ConceptualModel().hypothesize(causes(tutoring, score))
        model = cm.query(iv=tutoring, dv=score)

        assert isinstance(model, StatisticalModel)
        assert model.iv == "Tutoring"
        assert model.dv == "Score"
        assert "Tutoring" in model.main_effects

    def test_with_confounder(self):
        student = Unit("student", cardinality=100)
        ses = categories(student, "SES", order=["lower", "middle", "upper"])
        tutoring = categories(student, "Tutoring", cardinality=2)
        score = continuous(student, "Score")

        cm = (
            ConceptualModel()
            .assume(causes(ses, score))
            .assume(causes(ses, tutoring))
            .hypothesize(causes(tutoring, score))
        )
        model = cm.query(iv=tutoring, dv=score)
        assert "SES" in model.main_effects
        assert "Tutoring" in model.main_effects

    def test_formula_string(self):
        student = Unit("student", cardinality=100)
        tutoring = categories(student, "Tutoring", cardinality=2)
        score = continuous(student, "Score")

        cm = ConceptualModel().hypothesize(causes(tutoring, score))
        model = cm.query(iv=tutoring, dv=score)

        formula = model.formula()
        assert formula.startswith("Score ~ ")
        assert "Tutoring" in formula


class TestFamilyLink:
    def test_continuous_dv_family(self):
        student = Unit("student", cardinality=100)
        x = categories(student, "X", cardinality=2)
        y = continuous(student, "Y")

        cm = ConceptualModel().hypothesize(causes(x, y))
        model = cm.query(iv=x, dv=y)
        assert "Inverse Gaussian" in model.family_candidates
        assert "Gaussian" in model.family_candidates

    def test_counts_dv_family(self):
        student = Unit("student", cardinality=100)
        x = categories(student, "X", cardinality=2)
        y = counts(student, "Y")

        cm = ConceptualModel().hypothesize(causes(x, y))
        model = cm.query(iv=x, dv=y)
        assert "Poisson" in model.family_candidates

    def test_with_family_override(self):
        student = Unit("student", cardinality=100)
        x = categories(student, "X", cardinality=2)
        y = continuous(student, "Y")

        cm = ConceptualModel().hypothesize(causes(x, y))
        model = cm.query(iv=x, dv=y)
        model2 = model.with_family("Gamma", link="log")
        assert model2.family == "Gamma"
        assert model2.link == "log"


class TestNestingInFormula:
    def test_nesting_random_intercept_in_formula(self):
        classroom = Unit("classroom", cardinality=10)
        student = Unit("student", cardinality=100, nests_within=classroom)
        tutoring = categories(student, "Tutoring", cardinality=2)
        score = continuous(student, "Score")

        cm = ConceptualModel().hypothesize(causes(tutoring, score))
        model = cm.query(iv=tutoring, dv=score)
        formula = model.formula()
        assert "(1 | classroom)" in formula


class TestInteractionInFormula:
    def test_interaction_term(self):
        student = Unit("student", cardinality=100)
        ses = categories(student, "SES", order=["lower", "middle", "upper"])
        tutoring = categories(student, "Tutoring", cardinality=2)
        score = continuous(student, "Score")

        cm = (
            ConceptualModel()
            .assume(causes(ses, score))
            .assume(causes(ses, tutoring))
            .hypothesize(causes(tutoring, score))
            .interacts(ses, tutoring, dv=score)
        )
        model = cm.query(iv=tutoring, dv=score)
        assert "SES:Tutoring" in model.interaction_effects
        assert "SES:Tutoring" in model.formula()


class TestSummary:
    def test_summary_dict(self):
        student = Unit("student", cardinality=100)
        x = categories(student, "X", cardinality=2)
        y = continuous(student, "Y")

        cm = ConceptualModel().hypothesize(causes(x, y))
        model = cm.query(iv=x, dv=y)
        s = model.summary()
        assert s["dv"] == "Y"
        assert s["iv"] == "X"
        assert isinstance(s["family_candidates"], dict)
