import numpy as np
from scipy.signal import butter, sosfiltfilt


def bandpass_filter(
    voltage: np.ndarray, sample_rate: int, low_hz: float, high_hz: float
) -> np.ndarray:
    """
    Zero-phase Butterworth bandpass along the sample axis.
    voltage: (n_channels, n_samples)
    """
    sos = butter(
        4,
        [low_hz / (sample_rate / 2), high_hz / (sample_rate / 2)],
        btype="band",
        output="sos",
    )
    return sosfiltfilt(sos, voltage, axis=1).astype(np.float32)


def extract_snippets(
    spike_band: np.ndarray, positions: np.ndarray, length: int
) -> np.ndarray:
    """
    Extract fixed-length snippets centred on the given positions.
    Pads with zeros at the edges.
    Returns: (n_spikes, n_channels, length)
    """
    half = length // 2
    n_channels, n_samples = spike_band.shape
    padded = np.pad(spike_band, ((0, 0), (half, half)), mode="constant")
    out = np.empty((len(positions), n_channels, length), dtype=np.float32)
    for i, p in enumerate(positions):
        out[i] = padded[:, p : p + length]
    return out