"""Public helper module.

Conceptually, this module groups the small signal-processing and binning
helpers used by decoding, PSTH, and subspace pipelines.

It exists as a separate unit so window indexing, single-event PSTH counting,
and temporal aggregation remain consistent across higher-level modules.

It connects raw time-aligned signals and spike times to the compact arrays
used by downstream analyses.
"""

from __future__ import annotations

from typing import Optional, Tuple

import numpy as np

__all__ = [
    "concat_rows",
    "moving_sum_windows",
    "nearest_bin",
    "psth_simple_counts_single_event",
]


def nearest_bin(x: np.ndarray, t: float) -> int:
    """Return the index of the time bin nearest to `t`."""

    return int(np.argmin(np.abs(np.asarray(x, dtype=np.float32) - float(t))))



def psth_simple_counts_single_event(
    spike_times: np.ndarray,
    event_time: float,
    pre_time: float,
    post_time: float,
    bin_size: float,
    bin_step: float,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Compute sliding-bin PSTH counts for one alignment event.

    Parameters
    ----------
    spike_times:
        One-dimensional spike-time vector for a single unit.
    event_time:
        Alignment timestamp.
    pre_time, post_time:
        Relative analysis window in seconds.
    bin_size:
        Counting-window size in seconds.
    bin_step:
        Step between successive windows in seconds.

    Returns
    -------
    tuple[np.ndarray, np.ndarray, np.ndarray]
        Rates, bin centers, and raw counts for the single event.
    """

    # Mirror MATLAB PSTH_Simple.m applied to a single anchor (Code_M spontaneous-
    # lick counting calls PSTH_Simple(SpikeTimes, lick_t, ...)). Build the
    # RelativeEdges = pre:step:post grid inclusively (like the MATLAB colon),
    # count with histcounts semantics (half-open bins, last bin closed), and label
    # window centers as RelativeEdges + WindowSize. Keep float64: spike times are
    # large absolute session timestamps (up to ~10^4 s) whose float32 rounding
    # (~0.5 ms) would shift spikes across bin edges relative to Code_M.
    relative_edges = np.arange(pre_time, post_time + 1e-12, bin_step, dtype=np.float64)
    n_bins = max(int(relative_edges.size - 1), 0)
    counts = np.zeros(n_bins, dtype=np.float64)
    centers = relative_edges[:n_bins] + float(bin_size)

    st = np.asarray(spike_times, dtype=np.float64).reshape(-1)
    if st.size == 0 or n_bins == 0:
        rates = counts / max(float(bin_size), 1e-12)
        return rates, centers, counts

    st = np.sort(st)
    bin_edges = float(event_time) + relative_edges
    idx = np.searchsorted(st, bin_edges, side="left")
    idx[-1] = np.searchsorted(st, bin_edges[-1], side="right")
    counts = np.diff(idx).astype(np.float64)

    rates = counts / max(float(bin_size), 1e-12)
    return rates, centers, counts



def moving_sum_windows(arr: np.ndarray, win: int, step: int) -> np.ndarray:
    """Compute sliding or blockwise sums along axis 0.

    Parameters
    ----------
    arr:
        Input array whose first axis is interpreted as time.
    win:
        Window length in samples.
    step:
        Step between successive windows in samples.

    Returns
    -------
    np.ndarray
        Summed windows with the same trailing dimensions as `arr`.
    """

    # MATLAB movsum_pg.m operates in double precision; keep float64 so the
    # re-binned spike sums match Code_M exactly.
    arr = np.asarray(arr, dtype=np.float64)
    n_time = arr.shape[0]

    if win <= 0 or step <= 0 or n_time < win:
        return np.zeros((0, *arr.shape[1:]), dtype=np.float64)

    if win == step:
        # Use a reshape-based fast path for non-overlapping windows.
        n_blocks = n_time // step
        n_trim = n_blocks * step
        if n_trim == 0:
            return np.zeros((0, *arr.shape[1:]), dtype=np.float64)
        return arr[:n_trim].reshape(n_blocks, step, *arr.shape[1:]).sum(axis=1, dtype=np.float64)

    starts = list(range(0, n_time - win + 1, step))
    out = np.empty((len(starts), *arr.shape[1:]), dtype=np.float64)
    for i, start in enumerate(starts):
        out[i] = arr[start : start + win].sum(axis=0, dtype=np.float64)
    return out



def concat_rows(accum: Optional[np.ndarray], block: np.ndarray) -> np.ndarray:
    """Append row blocks while preserving a 2D float32 matrix layout."""

    blk = np.asarray(block, dtype=np.float32)
    if blk.ndim == 1:
        blk = blk.reshape(1, -1)
    if accum is None:
        return blk
    return np.vstack([accum, blk])