"""TRF system simulation helpers."""

from __future__ import annotations

import numpy as np

from neuroweave.design.lag import make_lagged_matrix


def generate_gaussian_kernel(
    n_features: int,
    n_lags: int,
    center: float,
    width: float,
) -> np.ndarray:
    """Generate a Gaussian-shaped TRF kernel across lags.

    Parameters
    ----------
    n_features : int
        Number of stimulus features.
    n_lags : int
        Number of lag positions.
    center : float
        Center of the Gaussian bump in lag-index coordinates.
    width : float
        Standard deviation of the Gaussian bump.

    Returns
    -------
    np.ndarray
        Kernel array with shape ``(n_features, n_lags)``.
    """

    lag_index = np.arange(n_lags, dtype=float)
    lag_profile = np.exp(-0.5 * ((lag_index - center) / width) ** 2)
    feature_weights = np.linspace(0.8, 1.2, n_features, dtype=float)
    return feature_weights[:, np.newaxis] * lag_profile[np.newaxis, :]


def simulate_trf_response(
    X: np.ndarray,
    kernel: np.ndarray,
    noise_std: float = 0.1,
) -> np.ndarray:
    """Simulate a TRF response from stimulus features and a kernel.

    Parameters
    ----------
    X : np.ndarray
        Stimulus matrix with shape ``(time, features)``.
    kernel : np.ndarray
        TRF kernel with shape ``(features, lags)``.
    noise_std : float, optional
        Standard deviation of additive Gaussian noise.

    Returns
    -------
    np.ndarray
        Simulated response vector with shape ``(time,)``.
    """

    X_array = np.asarray(X, dtype=float)
    kernel_array = np.asarray(kernel, dtype=float)

    if X_array.ndim != 2:
        raise ValueError("X must have shape (time, features).")
    if kernel_array.ndim != 2:
        raise ValueError("kernel must have shape (features, lags).")
    if X_array.shape[1] != kernel_array.shape[0]:
        raise ValueError("Feature dimension of X and kernel must match.")

    min_lag = -(kernel_array.shape[1] // 2)
    max_lag = min_lag + kernel_array.shape[1] - 1
    X_lagged = make_lagged_matrix(X_array, (min_lag, max_lag))
    beta = kernel_array.T.reshape(-1)
    signal = X_lagged @ beta
    noise = np.random.normal(loc=0.0, scale=noise_std, size=X_array.shape[0])
    return signal + noise
