"""Lag design matrix utilities for encoding models."""

from __future__ import annotations

import numpy as np


def make_lagged_matrix(
    X: np.ndarray,
    lags: tuple[int, int],
) -> np.ndarray:
    """Construct a zero-padded lagged design matrix.

    Parameters
    ----------
    X : np.ndarray
        Input feature matrix with shape ``(time, features)``.
    lags : tuple[int, int]
        Inclusive lag range ``(min_lag, max_lag)``.

    Returns
    -------
    np.ndarray
        Lagged feature matrix with shape
        ``(time, features * number_of_lags)``.
    """

    X_array = np.asarray(X, dtype=float)
    if X_array.ndim == 1:
        X_array = X_array[:, np.newaxis]
    if X_array.ndim != 2:
        raise ValueError("X must be a 1D or 2D array.")

    min_lag, max_lag = lags
    if min_lag > max_lag:
        raise ValueError("lags must satisfy min_lag <= max_lag.")

    n_times, n_features = X_array.shape
    lag_values = range(min_lag, max_lag + 1)
    lagged_blocks: list[np.ndarray] = []

    for lag in lag_values:
        block = np.zeros((n_times, n_features), dtype=X_array.dtype)

        if lag < 0:
            block[-lag:, :] = X_array[: n_times + lag, :]
        elif lag > 0:
            block[: n_times - lag, :] = X_array[lag:, :]
        else:
            block[:, :] = X_array

        lagged_blocks.append(block)

    return np.hstack(lagged_blocks)
