"""ConceptualModel: the central object for declaring domain knowledge.

Ports ConceptualModel, assume(), hypothesize() from rTisane.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

import networkx as nx

from .variables import AbstractVariable, Measure, Unit, Time
from .relationships import Causes, Relates, Nests, Interacts
from .graph import build_graph

if TYPE_CHECKING:
    from .statistical_model import StatisticalModel


# ---------------------------------------------------------------------------
# Wrappers that label a relationship as assumed vs hypothesized
# ---------------------------------------------------------------------------

@dataclass
class Assumption:
    """A relationship the analyst assumes (prior work / strong belief)."""
    relationship: Causes | Relates | Interacts


@dataclass
class Hypothesis:
    """A relationship the analyst hypothesises (focus of analysis)."""
    relationship: Causes | Relates | Interacts


# ---------------------------------------------------------------------------
# ConceptualModel
# ---------------------------------------------------------------------------

class ConceptualModel:
    """Graph-backed collection of variables and labeled relationships."""

    def __init__(self):
        self.variables: dict[str, AbstractVariable] = {}
        self.relationships: list[Assumption | Hypothesis] = []
        self.graph: nx.DiGraph = nx.DiGraph()
        self.interactions: list[Interacts] = []

    # -- helpers ----------------------------------------------------------

    def _track_variables(self, rel: Causes | Relates | Nests | Interacts):
        """Register variables mentioned in a relationship."""
        if isinstance(rel, Causes):
            self._add_var(rel.cause)
            self._add_var(rel.effect)
        elif isinstance(rel, Relates):
            self._add_var(rel.lhs)
            self._add_var(rel.rhs)
        elif isinstance(rel, Nests):
            self._add_var(rel.base)
            self._add_var(rel.group)
        elif isinstance(rel, Interacts):
            for v in rel.variables:
                self._add_var(v)
            if rel.dv is not None:
                self._add_var(rel.dv)

    def _add_var(self, var: AbstractVariable):
        if var.name not in self.variables:
            self.variables[var.name] = var

    def _rebuild_graph(self):
        self.graph = build_graph(self.relationships)

    # -- public API -------------------------------------------------------

    def assume(self, relationship: Causes | Relates) -> "ConceptualModel":
        """Add an assumed relationship. Returns *self* for chaining."""
        a = Assumption(relationship)
        self.relationships.append(a)
        self._track_variables(relationship)
        self._rebuild_graph()
        return self

    def hypothesize(self, relationship: Causes | Relates) -> "ConceptualModel":
        """Add a hypothesised relationship. Returns *self* for chaining."""
        h = Hypothesis(relationship)
        self.relationships.append(h)
        self._track_variables(relationship)
        self._rebuild_graph()
        return self

    def interacts(self, *variables: AbstractVariable, dv: AbstractVariable) -> "ConceptualModel":
        """Annotate an interaction among *variables* w.r.t. *dv*."""
        interaction = Interacts(variables=list(variables), dv=dv)
        self.interactions.append(interaction)
        self._track_variables(interaction)
        return self

    # -- query ------------------------------------------------------------

    def query(self, *, iv: AbstractVariable, dv: AbstractVariable, data=None) -> "StatisticalModel":
        """Infer a statistical model for the effect of *iv* on *dv*.

        Requires at least one hypothesised relationship involving *iv* and *dv*.
        """
        from .inference.confounders import infer_confounders
        from .inference.random_effects import infer_random_effects
        from .inference.family_link import infer_family_link_functions
        from .statistical_model import StatisticalModel

        # Validate: must have at least one hypothesis
        hypotheses = [r for r in self.relationships if isinstance(r, Hypothesis)]
        if not hypotheses:
            raise ValueError("ConceptualModel must contain at least one hypothesized relationship to query.")

        # Validate: hypothesis must connect iv and dv
        has_iv_dv = False
        for h in hypotheses:
            rel = h.relationship
            if isinstance(rel, Causes):
                if rel.cause.name == iv.name and rel.effect.name == dv.name:
                    has_iv_dv = True
            elif isinstance(rel, Relates):
                names = {rel.lhs.name, rel.rhs.name}
                if iv.name in names and dv.name in names:
                    has_iv_dv = True
        if not has_iv_dv:
            raise ValueError(
                f"No hypothesis connects iv={iv.name!r} and dv={dv.name!r}."
            )

        # 1. Confounders
        confounders = infer_confounders(self, iv, dv)

        # 2. Random effects
        relevant_interactions = [
            ix for ix in self.interactions if ix.dv is not None and ix.dv.name == dv.name
        ]
        random_effects = infer_random_effects(
            confounders=confounders,
            interactions=relevant_interactions,
            conceptual_model=self,
            iv=iv,
            dv=dv,
        )

        # 3. Family / link
        family_candidates = infer_family_link_functions(dv)

        # 4. Interaction effects
        interaction_effects: list[str] = []
        for ix in relevant_interactions:
            var_names = [v.name for v in ix.variables]
            interaction_effects.append(":".join(var_names))

        # 5. Build StatisticalModel
        first_family = next(iter(family_candidates))
        first_link = family_candidates[first_family][0] if family_candidates[first_family] else None

        return StatisticalModel(
            dv=dv.name,
            iv=iv.name,
            main_effects=[iv.name] + confounders,
            interaction_effects=interaction_effects,
            random_effects=random_effects,
            family_candidates=family_candidates,
            family=first_family,
            link=first_link,
        )
