"""Random-effects inference from data-measurement relationships.

Implements Barr et al. (2013) "Keep it maximal" and Barr (2013)
"Random effects structure for testing interactions."

Core rules:
1. Always include a random intercept for DV's unit when DV has repeated measures.
2. For each within-unit main effect (IV + confounders): add a random slope by unit.
3. For interactions: if all factors are within-unit, add a random slope for the
   interaction term; if mixed, add random slopes for the within-unit subset.
4. Group all random effects sharing the same grouping unit into a single term
   so lme4 estimates correlations between slopes (Barr's default recommendation).
5. Add random intercepts for nesting parents.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from ..variables import Measure, Time, Unit, Per, Exactly, AtMost

if TYPE_CHECKING:
    from ..conceptual_model import ConceptualModel
    from ..variables import AbstractVariable
    from ..relationships import Interacts


# ---------------------------------------------------------------------------
# Random-effect types
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class RandomIntercept:
    """Random intercept grouped by *group*."""
    group: Unit | Time

    def __str__(self):
        return f"(1 | {self.group.name})"


@dataclass(frozen=True)
class RandomSlope:
    """Random slope for *variable* grouped by *group*."""
    variable: Measure
    group: Unit

    def __str__(self):
        return f"({self.variable.name} | {self.group.name})"


@dataclass(frozen=True)
class RandomSlopeInteraction:
    """Random slope for an interaction term grouped by *group*."""
    variables: tuple[str, ...]
    group: Unit

    @property
    def term(self) -> str:
        return ":".join(self.variables)

    def __str__(self):
        return f"({self.term} | {self.group.name})"


RandomEffect = RandomIntercept | RandomSlope | RandomSlopeInteraction


@dataclass
class RandomEffectGroup:
    """A grouped random-effect term: ``(1 + x + z | group)``."""
    group: Unit | Time
    intercept: bool = True
    slopes: list[str] = field(default_factory=list)

    def __str__(self):
        parts: list[str] = []
        if self.intercept:
            parts.append("1")
        if not self.intercept and not self.slopes:
            parts.append("1")
        parts.extend(self.slopes)
        return f"({' + '.join(parts)} | {self.group.name})"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _num_value(n) -> int:
    """Extract a plain int from a NumberOfInstances value."""
    if isinstance(n, int):
        return n
    if isinstance(n, (Exactly, AtMost)):
        return n.value
    return 0  # Per or unknown — handled separately


def _gt_one(n) -> bool:
    """Check if a number-of-instances value is > 1."""
    if isinstance(n, Per):
        return True  # Per always implies repeated measures
    return _num_value(n) > 1


def _eq_one(n) -> bool:
    """Check if a number-of-instances value equals 1."""
    if isinstance(n, Per):
        return False
    return _num_value(n) == 1


def _is_within_unit(var: "AbstractVariable") -> bool:
    """True if the variable is measured multiple times per unit (within-subjects)."""
    if not isinstance(var, Measure):
        return False
    return _gt_one(var.number_of_instances)


def _is_between_unit(var: "AbstractVariable") -> bool:
    """True if the variable is measured once per unit (between-subjects)."""
    if not isinstance(var, Measure):
        return False
    return _eq_one(var.number_of_instances)


# ---------------------------------------------------------------------------
# Grouping: merge individual effects into grouped terms
# ---------------------------------------------------------------------------

def group_random_effects(effects: list[RandomEffect]) -> list[RandomEffectGroup]:
    """Merge random effects sharing the same grouping unit into grouped terms.

    Produces ``(1 + x + z | group)`` style terms for lme4.
    """
    from collections import OrderedDict

    groups: OrderedDict[str, RandomEffectGroup] = OrderedDict()

    for e in effects:
        if isinstance(e, RandomIntercept):
            key = e.group.name
            if key not in groups:
                groups[key] = RandomEffectGroup(group=e.group, intercept=True)
            else:
                groups[key].intercept = True
        elif isinstance(e, RandomSlope):
            key = e.group.name
            if key not in groups:
                groups[key] = RandomEffectGroup(group=e.group, intercept=False)
            slope_name = e.variable.name
            if slope_name not in groups[key].slopes:
                groups[key].slopes.append(slope_name)
        elif isinstance(e, RandomSlopeInteraction):
            key = e.group.name
            if key not in groups:
                groups[key] = RandomEffectGroup(group=e.group, intercept=False)
            slope_name = e.term
            if slope_name not in groups[key].slopes:
                groups[key].slopes.append(slope_name)

    return list(groups.values())


# ---------------------------------------------------------------------------
# Main inference
# ---------------------------------------------------------------------------

def infer_random_effects(
    confounders: list[str],
    interactions: list["Interacts"],
    conceptual_model: "ConceptualModel",
    iv: "AbstractVariable",
    dv: "AbstractVariable",
) -> list[RandomEffectGroup]:
    """Infer random effects from measurement relationships.

    Returns a list of RandomEffectGroup objects (grouped by unit).
    """
    effects: list[RandomEffect] = []

    # Gather main-effect variables (iv + confounders)
    main_vars: list[AbstractVariable] = []
    main_vars.append(iv)
    for cname in confounders:
        var = conceptual_model.variables.get(cname)
        if var is not None:
            main_vars.append(var)

    # Track within-unit slopes for reuse by crossed units
    within_unit_slopes: list[Measure] = []

    # -- Repeated measures ------------------------------------------------
    if isinstance(dv, Measure) and _gt_one(dv.number_of_instances):
        # Rule 1: Always add random intercept for DV's unit
        effects.append(RandomIntercept(group=dv.unit))

        # Rule 2: Random slopes for within-unit main effects
        for var in main_vars:
            if isinstance(var, Measure):
                if _is_within_unit(var):
                    effects.append(RandomSlope(variable=var, group=dv.unit))
                    within_unit_slopes.append(var)
                elif _is_between_unit(var):
                    # Between-subjects: no random slope needed,
                    # the random intercept already covers it.
                    # If the variable's unit differs from DV's unit, add
                    # a random intercept for that unit.
                    if var.unit.name != dv.unit.name:
                        effects.append(RandomIntercept(group=var.unit))

        # Rule 3: Random slopes for interactions
        for interaction in interactions:
            within_vars: list[str] = []
            for var in interaction.variables:
                if isinstance(var, Measure) and _is_within_unit(var):
                    within_vars.append(var.name)

            if len(within_vars) == len(interaction.variables):
                # All within-unit: add random slope for the full interaction
                effects.append(RandomSlopeInteraction(
                    variables=tuple(v.name for v in interaction.variables),
                    group=dv.unit,
                ))
            elif within_vars:
                # Mixed: add random slopes for the within-unit subset
                # (following Barr 2013's partial recommendation)
                if len(within_vars) > 1:
                    effects.append(RandomSlopeInteraction(
                        variables=tuple(within_vars),
                        group=dv.unit,
                    ))
                # Individual within-unit slopes are already added above in Rule 2

        # -- Crossed random effects (auto-detected from Per) --------------
        # When DV is measured Per(n, unit) where the Per variable is a Unit,
        # that unit is a crossed grouping factor (e.g., items in a
        # subjects-by-items design).  Apply Barr et al.'s maximal heuristic:
        # intercept + within-subject slopes for the crossed unit.
        if isinstance(dv.number_of_instances, Per):
            dv_per = dv.number_of_instances
            per_var = dv_per.variable
            if isinstance(per_var, Unit):
                crossed_unit = per_var
                effects.append(RandomIntercept(group=crossed_unit))
                for var in within_unit_slopes:
                    effects.append(RandomSlope(variable=var, group=crossed_unit))
                for interaction in interactions:
                    ix_within: list[str] = []
                    for var in interaction.variables:
                        if isinstance(var, Measure) and _is_within_unit(var):
                            ix_within.append(var.name)
                    if len(ix_within) == len(interaction.variables):
                        effects.append(RandomSlopeInteraction(
                            variables=tuple(v.name for v in interaction.variables),
                            group=crossed_unit,
                        ))
            elif isinstance(per_var, Measure):
                effects.append(RandomIntercept(group=per_var.unit))
            elif isinstance(per_var, Time):
                effects.append(RandomIntercept(group=per_var))

    # -- Nesting ----------------------------------------------------------
    if isinstance(dv, Measure):
        dv_unit = dv.unit
        if dv_unit.nests_within is not None:
            if isinstance(dv_unit.nests_within, list):
                for parent_unit in dv_unit.nests_within:
                    effects.append(RandomIntercept(group=parent_unit))
            else:
                effects.append(RandomIntercept(group=dv_unit.nests_within))

    # Group by unit and deduplicate
    return group_random_effects(effects)
