"""Design-matrix builder for declarative model specifications."""

from __future__ import annotations

from typing import Iterable

import numpy as np

from neuroweave.data.dataset import Dataset
from neuroweave.design.design_matrix import DesignMatrix
from neuroweave.design.lag import make_lagged_matrix
from neuroweave.design.lag_metadata import LagMetadata

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


def build_design_matrix(
    dataset: Dataset,
    terms: Iterable[FeatureTerm],
    lags: tuple[int, int],
    interactions: Iterable[InteractionTerm] | None = None,
    transforms: dict[str, FeatureTransform | list[FeatureTransform]] | None = None,
) -> DesignMatrix:
    """Build a lagged design matrix from a dataset and model terms.

    Parameters
    ----------
    dataset : Dataset
        Dataset providing neural data and existing stimulus features.
    terms : Iterable[FeatureTerm]
        Predictor terms to extract.
    lags : tuple[int, int]
        Inclusive lag range ``(min_lag, max_lag)``.
    interactions : Iterable[InteractionTerm] | None, optional
        Optional multiplicative interaction terms.
    transforms : dict[str, FeatureTransform | list[FeatureTransform]] | None, optional
        Optional per-predictor transforms keyed by predictor name.

    Returns
    -------
    DesignMatrix
        Lag-expanded design matrix.
    """

    transform_map = transforms or {}
    feature_columns: list[np.ndarray] = []
    feature_names: list[str] = []
    transformed_lookup: dict[str, tuple[np.ndarray, str]] = {}

    for term in terms:
        column = term.apply(dataset)
        output_name = term.name

        configured = transform_map.get(term.name, [])
        if isinstance(configured, FeatureTransform):
            configured_transforms = [configured]
        else:
            configured_transforms = list(configured)

        for transform in configured_transforms:
            column = transform.apply(column)
            output_name = transform.output_name(output_name)

        feature_columns.append(column.reshape(-1))
        feature_names.append(output_name)
        transformed_lookup[term.name] = (column.reshape(-1), output_name)

    if interactions is not None:
        for interaction in interactions:
            if interaction.left not in transformed_lookup or interaction.right not in transformed_lookup:
                raise ValueError("Interaction features must be included in predictors.")
            left_column, left_name = transformed_lookup[interaction.left]
            right_column, right_name = transformed_lookup[interaction.right]
            feature_columns.append(left_column * right_column)
            feature_names.append(f"{left_name} × {right_name}")

    if not feature_columns:
        raise ValueError("At least one predictor is required to build a design matrix.")

    stacked = np.column_stack(feature_columns)
    lagged = make_lagged_matrix(stacked, lags)
    metadata = LagMetadata(
        n_features=stacked.shape[1],
        lags=lags,
        feature_names=feature_names,
    )
    return DesignMatrix(X=lagged, metadata=metadata, sampling_rate=dataset.sampling_rate)
