"""Tests for inference/random_effects.py."""

from rtisanepy import Unit, continuous, categories, causes, ConceptualModel, Per, Exactly
from rtisanepy.relationships import Interacts
from rtisanepy.inference.random_effects import (
    RandomIntercept,
    RandomSlope,
    RandomSlopeInteraction,
    RandomEffectGroup,
    infer_random_effects,
    group_random_effects,
)


# ---------------------------------------------------------------------------
# Helper to find a group by unit name
# ---------------------------------------------------------------------------

def _find_group(groups: list[RandomEffectGroup], unit_name: str) -> RandomEffectGroup | None:
    for g in groups:
        if g.group.name == unit_name:
            return g
    return None


# ---------------------------------------------------------------------------
# Existing tests (updated for RandomEffectGroup return type)
# ---------------------------------------------------------------------------

class TestBetweenSubjects:
    def test_between_subjects_gives_random_intercept(self):
        """DV has >1 instances, IV has 1 instance → random intercept by DV's unit."""
        student = Unit("student", cardinality=100)
        score = continuous(student, "Score", number_of_instances=3)
        tutoring = categories(student, "Tutoring", cardinality=2)

        cm = ConceptualModel().hypothesize(causes(tutoring, score))

        groups = infer_random_effects(
            confounders=[], interactions=[], conceptual_model=cm,
            iv=tutoring, dv=score,
        )
        g = _find_group(groups, "student")
        assert g is not None
        assert g.intercept is True
        # Between-subjects IV should NOT get a random slope
        assert "Tutoring" not in g.slopes


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

        cm = ConceptualModel().hypothesize(causes(tutoring, score))
        groups = infer_random_effects(
            confounders=[], interactions=[], conceptual_model=cm,
            iv=tutoring, dv=score,
        )
        g = _find_group(groups, "classroom")
        assert g is not None
        assert g.intercept is True

    def test_nesting_with_within_subjects_iv(self):
        """Nested design with within-subjects IV gets both nesting intercept and slope."""
        classroom = Unit("classroom", cardinality=10)
        student = Unit("student", cardinality=100, nests_within=classroom)
        condition = categories(student, "Condition", cardinality=3, number_of_instances=3)
        score = continuous(student, "Score", number_of_instances=3)

        cm = ConceptualModel().hypothesize(causes(condition, score))
        groups = infer_random_effects(
            confounders=[], interactions=[], conceptual_model=cm,
            iv=condition, dv=score,
        )
        # Random intercept + slope for student
        student_g = _find_group(groups, "student")
        assert student_g is not None
        assert student_g.intercept is True
        assert "Condition" in student_g.slopes

        # Random intercept for classroom (nesting parent)
        classroom_g = _find_group(groups, "classroom")
        assert classroom_g is not None
        assert classroom_g.intercept is True


class TestWithinSubjects:
    def test_within_subjects_per(self):
        """IV measured Per(1, Time) → random intercept by unit."""
        student = Unit("student", cardinality=100)
        week = categories(student, "Week", cardinality=4)
        score = continuous(student, "Score", number_of_instances=Per(number=1, variable=week))
        tutoring = categories(student, "Tutoring", cardinality=2)

        cm = ConceptualModel().hypothesize(causes(tutoring, score))
        groups = infer_random_effects(
            confounders=[], interactions=[], conceptual_model=cm,
            iv=tutoring, dv=score,
        )
        g = _find_group(groups, "student")
        assert g is not None
        assert g.intercept is True


# ---------------------------------------------------------------------------
# Case 1: Within-subjects IV gets random slope
# ---------------------------------------------------------------------------

class TestWithinSubjectsRandomSlope:
    def test_within_subjects_iv_integer_gets_random_slope(self):
        """Within-subjects IV (number_of_instances=3) → random slope by unit.

        This is the core recommendation of Barr et al. (2013).
        Expected formula: Score ~ Condition + (1 + Condition | participant)
        """
        participant = Unit("participant", cardinality=40)
        condition = categories(participant, "Condition", cardinality=3,
                               number_of_instances=3)
        rt = continuous(participant, "RT", number_of_instances=3)

        cm = ConceptualModel().hypothesize(causes(condition, rt))
        groups = infer_random_effects(
            confounders=[], interactions=[], conceptual_model=cm,
            iv=condition, dv=rt,
        )
        g = _find_group(groups, "participant")
        assert g is not None
        assert g.intercept is True
        assert "Condition" in g.slopes

    def test_within_subjects_per_iv_gets_random_slope(self):
        """IV with Per() number_of_instances also gets a random slope."""
        student = Unit("student", cardinality=50)
        week = categories(student, "Week", cardinality=4)
        mood = categories(student, "Mood", cardinality=3,
                          number_of_instances=Per(number=1, variable=week))
        score = continuous(student, "Score",
                           number_of_instances=Per(number=1, variable=week))

        cm = ConceptualModel().hypothesize(causes(mood, score))
        groups = infer_random_effects(
            confounders=[], interactions=[], conceptual_model=cm,
            iv=mood, dv=score,
        )
        g = _find_group(groups, "student")
        assert g is not None
        assert "Mood" in g.slopes

    def test_between_subjects_iv_no_slope(self):
        """Between-subjects IV should NOT get a random slope."""
        participant = Unit("participant", cardinality=40)
        group = categories(participant, "Group", cardinality=2)
        score = continuous(participant, "Score", number_of_instances=5)

        cm = ConceptualModel().hypothesize(causes(group, score))
        groups = infer_random_effects(
            confounders=[], interactions=[], conceptual_model=cm,
            iv=group, dv=score,
        )
        g = _find_group(groups, "participant")
        assert g is not None
        assert "Group" not in g.slopes


# ---------------------------------------------------------------------------
# Case 3: Random intercept always present for DV unit with repeated measures
# ---------------------------------------------------------------------------

class TestRandomInterceptAlwaysPresent:
    def test_intercept_with_only_within_subjects_vars(self):
        """Random intercept for DV unit should be present even when all IVs are within."""
        subj = Unit("subject", cardinality=50)
        cond = categories(subj, "Cond", cardinality=2, number_of_instances=2)
        score = continuous(subj, "Score", number_of_instances=2)

        cm = ConceptualModel().hypothesize(causes(cond, score))
        groups = infer_random_effects(
            confounders=[], interactions=[], conceptual_model=cm,
            iv=cond, dv=score,
        )
        g = _find_group(groups, "subject")
        assert g is not None
        assert g.intercept is True

    def test_no_repeated_measures_no_random_intercept(self):
        """When DV has number_of_instances=1, no random intercept needed."""
        subj = Unit("subject", cardinality=50)
        group_var = categories(subj, "Group", cardinality=2)
        score = continuous(subj, "Score")  # default: number_of_instances=1

        cm = ConceptualModel().hypothesize(causes(group_var, score))
        groups = infer_random_effects(
            confounders=[], interactions=[], conceptual_model=cm,
            iv=group_var, dv=score,
        )
        # No repeated measures → no random effects from this path
        # (but nesting could still add them)
        g = _find_group(groups, "subject")
        assert g is None


# ---------------------------------------------------------------------------
# Case 7: Within-unit confounders get random slopes
# ---------------------------------------------------------------------------

class TestConfounderRandomSlopes:
    def test_within_unit_confounder_gets_slope(self):
        """A within-unit confounder should get a random slope."""
        subj = Unit("subject", cardinality=50)
        fatigue = continuous(subj, "Fatigue", number_of_instances=3)
        condition = categories(subj, "Condition", cardinality=2,
                               number_of_instances=3)
        score = continuous(subj, "Score", number_of_instances=3)

        cm = (
            ConceptualModel()
            .assume(causes(fatigue, score))
            .assume(causes(fatigue, condition))
            .hypothesize(causes(condition, score))
        )
        model = cm.query(iv=condition, dv=score)
        formula = model.formula()
        # Fatigue is a within-unit confounder; should be a slope
        assert "Fatigue" in formula

        groups = model.random_effects
        g = _find_group(groups, "subject")
        assert g is not None
        assert "Fatigue" in g.slopes
        assert "Condition" in g.slopes

    def test_between_unit_confounder_no_slope(self):
        """A between-unit confounder should NOT get a random slope."""
        subj = Unit("subject", cardinality=50)
        gender = categories(subj, "Gender", cardinality=2)  # between
        condition = categories(subj, "Condition", cardinality=2,
                               number_of_instances=3)
        score = continuous(subj, "Score", number_of_instances=3)

        cm = (
            ConceptualModel()
            .assume(causes(gender, score))
            .assume(causes(gender, condition))
            .hypothesize(causes(condition, score))
        )
        model = cm.query(iv=condition, dv=score)
        groups = model.random_effects
        g = _find_group(groups, "subject")
        assert g is not None
        assert "Gender" not in g.slopes


# ---------------------------------------------------------------------------
# Case 2: Interaction random slopes
# ---------------------------------------------------------------------------

class TestInteractionRandomSlopes:
    def test_all_within_interaction_gets_slope(self):
        """When all factors in an interaction are within-unit, add random slope
        for the full interaction term (Barr 2013)."""
        subj = Unit("subject", cardinality=50)
        a = categories(subj, "A", cardinality=2, number_of_instances=4)
        b = categories(subj, "B", cardinality=2, number_of_instances=4)
        y = continuous(subj, "Y", number_of_instances=4)

        interaction = Interacts(variables=[a, b], dv=y)

        groups = infer_random_effects(
            confounders=[], interactions=[interaction],
            conceptual_model=ConceptualModel()
                .assume(causes(b, y))
                .hypothesize(causes(a, y)),
            iv=a, dv=y,
        )
        g = _find_group(groups, "subject")
        assert g is not None
        assert g.intercept is True
        # Should have slopes for A, and for A:B interaction
        assert "A" in g.slopes
        assert "A:B" in g.slopes

    def test_mixed_interaction_no_full_slope(self):
        """When only some factors are within-unit, do NOT add a slope for the
        full interaction. Individual within-unit slopes are still added."""
        subj = Unit("subject", cardinality=50)
        a = categories(subj, "A", cardinality=2, number_of_instances=4)
        b = categories(subj, "B", cardinality=2)  # between-subjects
        y = continuous(subj, "Y", number_of_instances=4)

        interaction = Interacts(variables=[a, b], dv=y)

        groups = infer_random_effects(
            confounders=[], interactions=[interaction],
            conceptual_model=ConceptualModel()
                .assume(causes(b, y))
                .hypothesize(causes(a, y)),
            iv=a, dv=y,
        )
        g = _find_group(groups, "subject")
        assert g is not None
        assert "A" in g.slopes
        # Full interaction slope should NOT be present (mixed)
        assert "A:B" not in g.slopes

    def test_interactions_parameter_is_used(self):
        """Verify interactions parameter is actually processed (Case 2)."""
        subj = Unit("subject", cardinality=50)
        a = categories(subj, "A", cardinality=2, number_of_instances=4)
        b = categories(subj, "B", cardinality=2, number_of_instances=4)
        y = continuous(subj, "Y", number_of_instances=4)

        # Without interactions
        groups_no_ix = infer_random_effects(
            confounders=[], interactions=[],
            conceptual_model=ConceptualModel()
                .assume(causes(b, y))
                .hypothesize(causes(a, y)),
            iv=a, dv=y,
        )
        # With interactions
        interaction = Interacts(variables=[a, b], dv=y)
        groups_with_ix = infer_random_effects(
            confounders=[], interactions=[interaction],
            conceptual_model=ConceptualModel()
                .assume(causes(b, y))
                .hypothesize(causes(a, y)),
            iv=a, dv=y,
        )
        # The version with interactions should have the A:B slope
        g_no = _find_group(groups_no_ix, "subject")
        g_with = _find_group(groups_with_ix, "subject")
        assert "A:B" not in (g_no.slopes if g_no else [])
        assert "A:B" in (g_with.slopes if g_with else [])


# ---------------------------------------------------------------------------
# Case 5: Grouped formula output
# ---------------------------------------------------------------------------

class TestGroupedFormulaOutput:
    def test_single_group_combined_formula(self):
        """Multiple slopes for the same group produce a single (1 + x + z | group) term."""
        subj = Unit("subject", cardinality=50)
        a = categories(subj, "A", cardinality=2, number_of_instances=4)
        b = categories(subj, "B", cardinality=2, number_of_instances=4)
        y = continuous(subj, "Y", number_of_instances=4)

        cm = (
            ConceptualModel()
            .assume(causes(b, y))
            .hypothesize(causes(a, y))
        )
        model = cm.query(iv=a, dv=y)
        formula = model.formula()
        # Should have a single grouped term, not separate terms
        assert "(1 + A" in formula or "(1 + B" in formula
        # Should NOT have separate (A | subject) and (B | subject)
        assert "(A | subject)" not in formula
        assert "(B | subject)" not in formula

    def test_group_str_format(self):
        """RandomEffectGroup.__str__ produces correct format."""
        subj = Unit("subject", cardinality=50)
        g = RandomEffectGroup(group=subj, intercept=True, slopes=["A", "B"])
        assert str(g) == "(1 + A + B | subject)"

    def test_group_str_intercept_only(self):
        subj = Unit("subject", cardinality=50)
        g = RandomEffectGroup(group=subj, intercept=True, slopes=[])
        assert str(g) == "(1 | subject)"

    def test_group_random_effects_merges_correctly(self):
        """group_random_effects merges intercepts and slopes for same unit."""
        subj = Unit("subject", cardinality=50)
        effects = [
            RandomIntercept(group=subj),
            RandomSlope(variable=categories(subj, "A", cardinality=2), group=subj),
            RandomSlope(variable=categories(subj, "B", cardinality=2), group=subj),
        ]
        groups = group_random_effects(effects)
        assert len(groups) == 1
        g = groups[0]
        assert g.intercept is True
        assert "A" in g.slopes
        assert "B" in g.slopes

    def test_separate_groups_stay_separate(self):
        """Effects for different grouping units remain separate terms."""
        subj = Unit("subject", cardinality=50)
        classroom = Unit("classroom", cardinality=10)
        effects = [
            RandomIntercept(group=subj),
            RandomSlope(variable=categories(subj, "A", cardinality=2), group=subj),
            RandomIntercept(group=classroom),
        ]
        groups = group_random_effects(effects)
        assert len(groups) == 2
        subj_g = _find_group(groups, "subject")
        class_g = _find_group(groups, "classroom")
        assert subj_g is not None and class_g is not None
        assert subj_g.intercept is True
        assert "A" in subj_g.slopes
        assert class_g.intercept is True
        assert class_g.slopes == []


# ---------------------------------------------------------------------------
# Mixed design (within + between)
# ---------------------------------------------------------------------------

class TestMixedDesign:
    def test_mixed_design_formula(self):
        """Mixed design: within IV gets slope, between IV doesn't.

        Expected: Y ~ A + B + (1 + A | subject) where A is within, B is between.
        """
        subj = Unit("subject", cardinality=50)
        a = categories(subj, "A", cardinality=2, number_of_instances=3)
        b = categories(subj, "B", cardinality=2)  # between
        y = continuous(subj, "Y", number_of_instances=3)

        cm = (
            ConceptualModel()
            .assume(causes(b, y))
            .assume(causes(b, a))
            .hypothesize(causes(a, y))
        )
        model = cm.query(iv=a, dv=y)
        groups = model.random_effects
        g = _find_group(groups, "subject")
        assert g is not None
        assert g.intercept is True
        assert "A" in g.slopes
        assert "B" not in g.slopes


# ---------------------------------------------------------------------------
# Full formula integration
# ---------------------------------------------------------------------------

class TestFormulaIntegration:
    def test_within_subjects_full_formula(self):
        """Full formula for a within-subjects design."""
        participant = Unit("participant", cardinality=40)
        condition = categories(participant, "Condition", cardinality=3,
                               number_of_instances=3)
        rt = continuous(participant, "RT", number_of_instances=3)

        cm = ConceptualModel().hypothesize(causes(condition, rt))
        model = cm.query(iv=condition, dv=rt)
        formula = model.formula()
        assert "RT ~ Condition + (1 + Condition | participant)" == formula

    def test_within_subjects_interaction_full_formula(self):
        """Full formula for within-subjects 2x2 design with interaction."""
        subj = Unit("subject", cardinality=50)
        a = categories(subj, "A", cardinality=2, number_of_instances=4)
        b = categories(subj, "B", cardinality=2, number_of_instances=4)
        y = continuous(subj, "Y", number_of_instances=4)

        cm = (
            ConceptualModel()
            .assume(causes(b, y))
            .hypothesize(causes(a, y))
            .interacts(a, b, dv=y)
        )
        model = cm.query(iv=a, dv=y)
        formula = model.formula()
        # Should contain the interaction term and grouped random effects
        assert "A:B" in formula
        assert "(1 + A + B + A:B | subject)" in formula

    def test_nested_within_subjects_formula(self):
        """Nested design with within-subjects IV."""
        classroom = Unit("classroom", cardinality=10)
        student = Unit("student", cardinality=100, nests_within=classroom)
        condition = categories(student, "Condition", cardinality=2,
                               number_of_instances=2)
        score = continuous(student, "Score", number_of_instances=2)

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


# ---------------------------------------------------------------------------
# Crossed random effects (auto-inferred from Per(n, Unit) on DV)
# ---------------------------------------------------------------------------

class TestCrossedRandomEffects:
    def test_canonical_crossed_design(self):
        """Canonical psycholinguistics design: subjects x items.

        Per(1, item) on DV signals that item is a crossed grouping factor.
        Expected: RT ~ Condition + (1 + Condition | subject) + (1 + Condition | item)
        """
        subject = Unit("subject", cardinality=40)
        item = Unit("item", cardinality=40)
        condition = categories(subject, "Condition", cardinality=2,
                               number_of_instances=2)
        rt = continuous(subject, "RT", number_of_instances=Per(1, item))

        cm = ConceptualModel().hypothesize(causes(condition, rt))
        model = cm.query(iv=condition, dv=rt)
        formula = model.formula()

        assert "(1 + Condition | subject)" in formula
        assert "(1 + Condition | item)" in formula

    def test_crossed_design_full_formula(self):
        """Verify exact formula for crossed design."""
        subject = Unit("subject", cardinality=40)
        item = Unit("item", cardinality=40)
        condition = categories(subject, "Condition", cardinality=2,
                               number_of_instances=2)
        rt = continuous(subject, "RT", number_of_instances=Per(1, item))

        cm = ConceptualModel().hypothesize(causes(condition, rt))
        model = cm.query(iv=condition, dv=rt)
        assert model.formula() == (
            "RT ~ Condition + (1 + Condition | subject) + (1 + Condition | item)"
        )

    def test_crossed_intercept_only_when_between_subjects(self):
        """Crossed unit gets intercept but no slope for between-subjects IV."""
        subject = Unit("subject", cardinality=40)
        item = Unit("item", cardinality=40)
        group_var = categories(subject, "Group", cardinality=2)  # between
        rt = continuous(subject, "RT", number_of_instances=Per(1, item))

        cm = ConceptualModel().hypothesize(causes(group_var, rt))
        model = cm.query(iv=group_var, dv=rt)
        groups = model.random_effects
        item_g = _find_group(groups, "item")
        assert item_g is not None
        assert item_g.intercept is True
        # Between-subjects IV → no slope for item either
        assert "Group" not in item_g.slopes

    def test_crossed_with_interaction(self):
        """Crossed design with all-within interaction gets interaction slope
        for both grouping factors."""
        subject = Unit("subject", cardinality=40)
        item = Unit("item", cardinality=40)
        a = categories(subject, "A", cardinality=2, number_of_instances=4)
        b = categories(subject, "B", cardinality=2, number_of_instances=4)
        y = continuous(subject, "Y", number_of_instances=Per(1, item))

        cm = (
            ConceptualModel()
            .assume(causes(b, y))
            .hypothesize(causes(a, y))
            .interacts(a, b, dv=y)
        )
        model = cm.query(iv=a, dv=y)
        groups = model.random_effects

        subj_g = _find_group(groups, "subject")
        item_g = _find_group(groups, "item")
        assert subj_g is not None
        assert item_g is not None

        # Both groups should have A, B slopes and A:B interaction slope
        assert "A" in subj_g.slopes
        assert "B" in subj_g.slopes
        assert "A:B" in subj_g.slopes
        assert "A" in item_g.slopes
        assert "B" in item_g.slopes
        assert "A:B" in item_g.slopes

    def test_no_per_unit_no_crossed_effects(self):
        """Without Per(n, Unit) on DV, no crossed random effects are generated."""
        subject = Unit("subject", cardinality=40)
        condition = categories(subject, "Condition", cardinality=2,
                               number_of_instances=2)
        rt = continuous(subject, "RT", number_of_instances=2)

        cm = ConceptualModel().hypothesize(causes(condition, rt))
        model = cm.query(iv=condition, dv=rt)
        # Only subject group, no item group
        assert len(model.random_effects) == 1
        assert model.random_effects[0].group.name == "subject"

    def test_per_with_measure_not_unit_gives_intercept_only(self):
        """Per(1, Measure) adds intercept for measure's unit, not full crossing."""
        student = Unit("student", cardinality=100)
        week = categories(student, "Week", cardinality=4)
        score = continuous(student, "Score",
                           number_of_instances=Per(number=1, variable=week))
        tutoring = categories(student, "Tutoring", cardinality=2)

        cm = ConceptualModel().hypothesize(causes(tutoring, score))
        groups = infer_random_effects(
            confounders=[], interactions=[], conceptual_model=cm,
            iv=tutoring, dv=score,
        )
        # Week is a Measure, not a Unit — gets intercept for its unit, not slopes
        g = _find_group(groups, "student")
        assert g is not None
        assert g.intercept is True
