"""Kernel reshaping utilities for NeuroWeave."""

from __future__ import annotations

import numpy as np


def vector_to_kernel(
    beta: np.ndarray,
    n_features: int,
    lags: tuple[int, int],
) -> np.ndarray:
    """Convert a coefficient vector to a structured kernel.

    Parameters
    ----------
    beta : np.ndarray
        Regression coefficient vector with shape ``(n_features * n_lags,)``.
    n_features : int
        Number of stimulus features.
    lags : tuple[int, int]
        Inclusive lag range ``(min_lag, max_lag)``.

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

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

    n_lags = max_lag - min_lag + 1
    beta_array = np.asarray(beta, dtype=float).reshape(-1)
    expected_size = n_features * n_lags

    if beta_array.size != expected_size:
        raise ValueError(
            "beta size does not match n_features * number_of_lags."
        )

    return beta_array.reshape(n_features, n_lags)


def kernel_to_vector(kernel: np.ndarray) -> np.ndarray:
    """Convert a structured kernel to a coefficient vector.

    Parameters
    ----------
    kernel : np.ndarray
        Kernel array with shape ``(features, lags)``.

    Returns
    -------
    np.ndarray
        Regression coefficient vector with shape ``(features * lags,)``.
    """

    kernel_array = np.asarray(kernel, dtype=float)
    if kernel_array.ndim != 2:
        raise ValueError("kernel must be a 2D array.")
    return kernel_array.reshape(-1)


def lag_axis(lags: tuple[int, int]) -> np.ndarray:
    """Return the lag axis for a kernel.

    Parameters
    ----------
    lags : tuple[int, int]
        Inclusive lag range ``(min_lag, max_lag)``.

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

    min_lag, max_lag = lags
    if min_lag > max_lag:
        raise ValueError("lags must satisfy min_lag <= max_lag.")
    return np.arange(min_lag, max_lag + 1, dtype=int)


def split_kernel(kernel: np.ndarray) -> list[np.ndarray]:
    """Split a kernel into one lag profile per feature.

    Parameters
    ----------
    kernel : np.ndarray
        Kernel array with shape ``(features, lags)``.

    Returns
    -------
    list[np.ndarray]
        List of one-dimensional feature kernels with shape ``(lags,)``.
    """

    kernel_array = np.asarray(kernel, dtype=float)
    if kernel_array.ndim != 2:
        raise ValueError("kernel must be a 2D array.")
    return [kernel_array[index, :].copy() for index in range(kernel_array.shape[0])]


def kernel_peak_lag(kernel: np.ndarray, lags: tuple[int, int]) -> np.ndarray:
    """Find the peak lag of each feature kernel.

    Parameters
    ----------
    kernel : np.ndarray
        Kernel array with shape ``(features, lags)``.
    lags : tuple[int, int]
        Inclusive lag range ``(min_lag, max_lag)``.

    Returns
    -------
    np.ndarray
        Peak lag value for each feature.
    """

    kernel_array = np.asarray(kernel, dtype=float)
    if kernel_array.ndim != 2:
        raise ValueError("kernel must be a 2D array.")

    axis = lag_axis(lags)
    if kernel_array.shape[1] != axis.size:
        raise ValueError("kernel lag dimension does not match the provided lags.")

    peak_indices = np.argmax(kernel_array, axis=1)
    return axis[peak_indices]


def kernel_energy(kernel: np.ndarray) -> np.ndarray:
    """Compute the energy of each feature kernel.

    Parameters
    ----------
    kernel : np.ndarray
        Kernel array with shape ``(features, lags)``.

    Returns
    -------
    np.ndarray
        Energy per feature computed as the sum of squared lag weights.
    """

    kernel_array = np.asarray(kernel, dtype=float)
    if kernel_array.ndim != 2:
        raise ValueError("kernel must be a 2D array.")
    return np.sum(kernel_array**2, axis=1)
