"""Declarative model specification objects for NeuroWeave."""

from __future__ import annotations

from typing import Any

from neuroweave.data.dataset import Dataset
from neuroweave.design.design_matrix import DesignMatrix

from .builder import build_design_matrix
from .terms import FeatureTerm, InteractionTerm
from .transforms import FeatureTransform


class ModelSpec:
    """Declarative specification for an encoding model design.

    Parameters
    ----------
    predictors : list[str]
        Predictor feature names available in ``dataset.features``.
    lags : tuple[int, int]
        Inclusive lag window.
    interactions : list[tuple[str, str]] | None, optional
        Optional feature interactions.
    transforms : dict[str, FeatureTransform | list[FeatureTransform]] | None, optional
        Optional per-predictor transforms.
    """

    def __init__(
        self,
        predictors: list[str],
        lags: tuple[int, int],
        interactions: list[tuple[str, str]] | None = None,
        transforms: dict[str, FeatureTransform | list[FeatureTransform]] | None = None,
        groups: dict[str, list[str]] | None = None,
    ) -> None:
        self.predictors = list(predictors)
        self.lags = lags
        self.interactions = list(interactions) if interactions is not None else []
        self.transforms = dict(transforms) if transforms is not None else {}
        self.groups = dict(groups) if groups is not None else {}

    def build(self, dataset: Dataset) -> DesignMatrix:
        """Build a lagged design matrix from a dataset.

        Parameters
        ----------
        dataset : Dataset
            Dataset providing existing stimulus features.

        Returns
        -------
        DesignMatrix
            Lag-expanded design matrix for the specification.
        """

        missing = [name for name in self.predictors if name not in dataset.feature_names]
        if missing:
            raise ValueError(f"Missing predictors in dataset.features: {missing}")

        terms = [FeatureTerm(name) for name in self.predictors]
        interaction_terms = [InteractionTerm(left, right) for left, right in self.interactions]
        return build_design_matrix(
            dataset=dataset,
            terms=terms,
            lags=self.lags,
            interactions=interaction_terms,
            transforms=self.transforms,
        )

    def feature_names(self) -> list[str]:
        """Return predictor names prior to lag expansion."""

        names = list(self.predictors)
        names.extend([f"{left} × {right}" for left, right in self.interactions])
        return names

    def summary(self) -> str:
        """Return a readable model specification summary."""

        lines = [
            "ModelSpec",
            "---------------------------------",
            "Predictors:",
        ]
        lines.extend([f"  {name}" for name in self.predictors])
        lines.append("")
        lines.append("Interactions:")
        if self.interactions:
            lines.extend([f"  {left} × {right}" for left, right in self.interactions])
        else:
            lines.append("  None")
        lines.append("")
        lines.append("Groups:")
        if self.groups:
            lines.extend([f"  {name}: {', '.join(features)}" for name, features in self.groups.items()])
        else:
            lines.append("  None")
        lines.append("")
        lines.append(f"Lag window:\n  {self.lags}")
        lines.append("")
        lines.append(f"Total predictors (before lagging):\n  {len(self.feature_names())}")
        return "\n".join(lines)
