"""Design matrix container for NeuroWeave."""

from __future__ import annotations

import copy

import numpy as np

from neuroweave.design.lag_metadata import LagMetadata
from neuroweave.utils.kernel import vector_to_kernel


class DesignMatrix:
    """Container for a lagged design matrix and its metadata.

    Parameters
    ----------
    X : np.ndarray
        Design matrix with shape ``(time, coefficients)``.
    metadata : LagMetadata
        Lag metadata describing the matrix layout.
    sampling_rate : float | None, optional
        Sampling rate associated with the design matrix.
    """

    def __init__(
        self,
        X: np.ndarray,
        metadata: LagMetadata,
        sampling_rate: float | None = None,
    ) -> None:
        X_array = np.asarray(X, dtype=float)
        if X_array.ndim != 2:
            raise ValueError("X must be a 2D array.")
        if X_array.shape[1] != metadata.n_coefficients:
            raise ValueError(
                "X column count must match metadata.n_coefficients."
            )

        self.X = X_array.copy()
        self.metadata = metadata
        self.sampling_rate = sampling_rate

    @property
    def n_samples(self) -> int:
        """Number of time samples."""

        return self.X.shape[0]

    @property
    def n_features(self) -> int:
        """Number of stimulus features."""

        return self.metadata.n_features

    @property
    def n_lags(self) -> int:
        """Number of lag positions."""

        return self.metadata.n_lags

    @property
    def lag_axis(self) -> np.ndarray:
        """Lag axis associated with the design matrix."""

        return self.metadata.lag_axis

    def kernel_from_vector(self, beta: np.ndarray) -> np.ndarray:
        """Reconstruct a kernel from a flattened coefficient vector.

        Parameters
        ----------
        beta : np.ndarray
            Coefficient vector with shape ``(features * lags,)``.

        Returns
        -------
        np.ndarray
            Kernel with shape ``(features, lags)``.
        """

        return vector_to_kernel(
            beta=beta,
            n_features=self.n_features,
            lags=(self.metadata.min_lag, self.metadata.max_lag),
        )

    def feature_kernel(self, beta: np.ndarray, feature_index: int) -> np.ndarray:
        """Return the kernel for a single feature.

        Parameters
        ----------
        beta : np.ndarray
            Coefficient vector with shape ``(features * lags,)``.
        feature_index : int
            Zero-based feature index.

        Returns
        -------
        np.ndarray
            One-dimensional kernel for the selected feature.
        """

        if not 0 <= feature_index < self.n_features:
            raise ValueError("feature_index is out of bounds.")
        kernel = self.kernel_from_vector(beta)
        return kernel[feature_index, :].copy()

    def lag_slice(self, lag: int) -> np.ndarray:
        """Return matrix columns corresponding to a specific lag.

        Parameters
        ----------
        lag : int
            Lag value in the inclusive interval ``[min_lag, max_lag]``.

        Returns
        -------
        np.ndarray
            Design-matrix columns for the selected lag with shape
            ``(samples, features)``.
        """

        lag_index = self.metadata.lag_to_index(lag)
        column_indices = [
            self.metadata.coefficient_index(feature_index, lag_index)
            for feature_index in range(self.n_features)
        ]
        return self.X[:, column_indices].copy()

    def copy(self) -> DesignMatrix:
        """Return a deep copy of the design matrix container.

        Returns
        -------
        DesignMatrix
            Copied design matrix object.
        """

        return DesignMatrix(
            X=self.X.copy(),
            metadata=copy.deepcopy(self.metadata),
            sampling_rate=self.sampling_rate,
        )

    def __repr__(self) -> str:
        """Return a readable representation."""

        return (
            f"DesignMatrix(samples={self.n_samples}, "
            f"features={self.n_features}, "
            f"lags=({self.metadata.min_lag},{self.metadata.max_lag}))"
        )
