"""Lag metadata for structured design matrix configuration."""

from __future__ import annotations

from typing import Final

import numpy as np


class LagMetadata:
    """Immutable metadata describing a lagged design matrix.

    Parameters
    ----------
    n_features : int
        Number of stimulus features.
    lags : tuple[int, int]
        Inclusive lag range ``(min_lag, max_lag)``.
    feature_names : list[str] | None, optional
        Optional feature names. When omitted, default names are generated.
    """

    __slots__: Final = (
        "_feature_names",
        "_initialized",
        "_lag_axis",
        "_lags",
        "_n_features",
    )

    def __init__(
        self,
        n_features: int,
        lags: tuple[int, int],
        feature_names: list[str] | None = None,
    ) -> None:
        min_lag, max_lag = lags
        if max_lag < min_lag:
            raise ValueError("max_lag must be greater than or equal to min_lag.")
        if n_features < 1:
            raise ValueError("n_features must be at least 1.")

        if feature_names is None:
            names = [f"feature_{index}" for index in range(n_features)]
        else:
            if len(feature_names) != n_features:
                raise ValueError("feature_names length must match n_features.")
            names = list(feature_names)

        object.__setattr__(self, "_n_features", int(n_features))
        object.__setattr__(self, "_lags", (int(min_lag), int(max_lag)))
        object.__setattr__(
            self,
            "_lag_axis",
            np.arange(min_lag, max_lag + 1, dtype=int),
        )
        object.__setattr__(self, "_feature_names", tuple(names))
        object.__setattr__(self, "_initialized", True)

    def __setattr__(self, name: str, value: object) -> None:
        """Prevent mutation after initialization."""

        if getattr(self, "_initialized", False):
            raise AttributeError("LagMetadata is immutable.")
        object.__setattr__(self, name, value)

    @property
    def min_lag(self) -> int:
        """Minimum lag value."""

        return self._lags[0]

    @property
    def max_lag(self) -> int:
        """Maximum lag value."""

        return self._lags[1]

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

        return self._n_features

    @property
    def feature_names(self) -> list[str]:
        """Feature names."""

        return list(self._feature_names)

    @property
    def lag_axis(self) -> np.ndarray:
        """Lag axis spanning the inclusive lag interval.

        Returns
        -------
        np.ndarray
            Lag vector from ``min_lag`` to ``max_lag``.
        """

        return self._lag_axis.copy()

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

        Returns
        -------
        int
            Total number of lags.
        """

        return self.max_lag - self.min_lag + 1

    @property
    def n_coefficients(self) -> int:
        """Total number of regression coefficients.

        Returns
        -------
        int
            Product of features and lag positions.
        """

        return self.n_features * self.n_lags

    def kernel_shape(self) -> tuple[int, int]:
        """Return the canonical kernel shape.

        Returns
        -------
        tuple[int, int]
            Kernel shape ``(features, lags)``.
        """

        return (self.n_features, self.n_lags)

    def vector_shape(self) -> tuple[int]:
        """Return the flattened coefficient vector shape.

        Returns
        -------
        tuple[int]
            Vector shape ``(features * lags,)``.
        """

        return (self.n_coefficients,)

    def coefficient_index(self, feature_index: int, lag_index: int) -> int:
        """Map a feature/lag pair to the flattened coefficient index.

        Parameters
        ----------
        feature_index : int
            Zero-based feature index.
        lag_index : int
            Zero-based lag index.

        Returns
        -------
        int
            Position in the flattened regression vector.
        """

        if not 0 <= feature_index < self.n_features:
            raise ValueError("feature_index is out of bounds.")
        if not 0 <= lag_index < self.n_lags:
            raise ValueError("lag_index is out of bounds.")
        return feature_index * self.n_lags + lag_index

    def lag_to_index(self, lag: int) -> int:
        """Convert a lag value to its zero-based lag index.

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

        Returns
        -------
        int
            Zero-based lag index.
        """

        if lag < self.min_lag or lag > self.max_lag:
            raise ValueError("lag is outside the configured lag range.")
        return lag - self.min_lag

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

        return (
            f"LagMetadata(features={self.n_features}, "
            f"lags=({self.min_lag},{self.max_lag}), n_lags={self.n_lags})"
        )
