streaming_waveforms.py

"""
streaming_waveforms.py

Full reimplementation of:
- load_cupid_waveform
- CR–RC², triangular, trapezoid filters
- Baseline RMS estimators (MAD, P², t-digest)
- Triggering

Plus streaming-capable backends suitable for GUI / long files:
- CUPID binary reader in streaming mode
- Streaming filters
- Streaming RMS estimators
- Streaming trigger state machine
"""

from __future__ import annotations

import os
import struct
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Protocol, Union, Iterator

import numpy as np

# ----------------------------------------------------------------------
# CUPID waveform loader (batch) + streaming reader
# ----------------------------------------------------------------------


def load_cupid_waveform(
    path: str,
    seconds: Union[float, str] = "all",
    adc_units: bool = False,
    scale: Optional[float] = None,
):
    """
    Load CUPID waveform from binary file, preserving the DC level.

    Parameters
    ----------
    path : str
        Full path to .bin file.
    seconds : float or "all"
        Duration to load. If "all", loads the full run.
    adc_units : bool
        If True, return ADC counts (24-bit) instead of Volts.
    scale : float or None
        Optional multiplicative factor (e.g. convert to mV, A, etc.)

    Returns
    -------
    t : ndarray
        Time axis (seconds)
    y : ndarray
        Waveform in Volts (or ADC units, or scaled units)
    fs : float
        Sampling frequency
    fullscale : float
        ADC full-scale voltage from header
    """

    # === read header ===
    with open(path, "rb") as f:
        header = f.read(12)
    magic, fs, fullscale = struct.unpack("<Iff", header)

    # === number of samples ===
    filesize = os.path.getsize(path)
    bytes_data = filesize - 12
    total_samples = bytes_data // 4

    if seconds == "all":
        N = total_samples
    else:
        N = min(int(fs * float(seconds)), total_samples)

    # === load raw 32-bit samples in streaming blocks (but return full arrays) ===
    # block size in samples (tunable)
    block_size = 1_000_000

    y_chunks: List[np.ndarray] = []
    t_chunks: List[np.ndarray] = []

    remaining = N
    offset = 0

    with open(path, "rb") as f:
        f.seek(12)  # skip header

        while remaining > 0:
            count = min(block_size, remaining)
            data = np.fromfile(f, dtype="<u4", count=count)
            if data.size == 0:
                break

            # === decode 24-bit ADC data ===
            ADC_raw = (data >> 8) & 0xFFFFFF  # still uint32

            if adc_units:
                # return raw 24-bit ADC counts
                y_block = ADC_raw.astype(np.float64)
            else:
                # convert ADC → Volt avoiding uint32 underflow
                ADC_centered = ADC_raw.astype(np.int32) - 0x800000
                y_block = fullscale * ADC_centered / 0x800000

            # === optional physical scaling (e.g. -> mV or amps) ===
            if scale is not None:
                y_block = y_block * float(scale)

            # === time axis for this block ===
            t0 = offset / fs
            t_block = t0 + np.arange(data.size, dtype=np.float64) / fs

            y_chunks.append(y_block)
            t_chunks.append(t_block)

            remaining -= data.size
            offset += data.size

    if not y_chunks:
        return np.array([]), np.array([]), fs, fullscale

    y = np.concatenate(y_chunks)
    t = np.concatenate(t_chunks)

    return t, y, fs, fullscale


class CupidWaveformReader:
    """
    Streaming reader for CUPID binary waveforms.

    Usage
    -----
    reader = CupidWaveformReader(path, adc_units=False, scale=None, block_size=100000)
    for t_block, y_block in reader:
        ... process block ...

    Attributes
    ----------
    fs : float
        Sampling frequency [Hz]
    fullscale : float
        ADC full-scale voltage from header
    total_samples : int
        Total number of samples in the file
    """

    def __init__(
        self,
        path: str,
        *,
        adc_units: bool = False,
        scale: Optional[float] = None,
        block_size: int = 100_000,
    ) -> None:
        self.path = path
        self.adc_units = adc_units
        self.scale = None if scale is None else float(scale)
        self.block_size = int(block_size)

        self._file = open(path, "rb")

        header = self._file.read(12)
        self.magic, self.fs, self.fullscale = struct.unpack("<Iff", header)

        filesize = os.path.getsize(path)
        bytes_data = filesize - 12
        self.total_samples = bytes_data // 4

        self._read_samples = 0
        self._closed = False

    def __iter__(self) -> Iterator[tuple[np.ndarray, np.ndarray]]:
        return self

    def __next__(self) -> tuple[np.ndarray, np.ndarray]:
        if self._closed or self._read_samples >= self.total_samples:
            self.close()
            raise StopIteration

        remaining = self.total_samples - self._read_samples
        count = min(self.block_size, remaining)

        data = np.fromfile(self._file, dtype="<u4", count=count)
        if data.size == 0:
            self.close()
            raise StopIteration

        ADC_raw = (data >> 8) & 0xFFFFFF

        if self.adc_units:
            y = ADC_raw.astype(np.float64)
        else:
            ADC_centered = ADC_raw.astype(np.int32) - 0x800000
            y = self.fullscale * ADC_centered / 0x800000

        if self.scale is not None:
            y = y * self.scale

        t0 = self._read_samples / self.fs
        t = t0 + np.arange(data.size, dtype=np.float64) / self.fs

        self._read_samples += data.size
        return t, y

    def read_samples(self, n_samples: int) -> tuple[np.ndarray, np.ndarray]:
        """
        Read up to n_samples in one go.

        Returns
        -------
        t_block, y_block : np.ndarray
            Time axis and waveform block.
        """
        if self._closed or self._read_samples >= self.total_samples:
            return np.array([]), np.array([])

        remaining = self.total_samples - self._read_samples
        count = min(int(n_samples), remaining)

        data = np.fromfile(self._file, dtype="<u4", count=count)
        if data.size == 0:
            self.close()
            return np.array([]), np.array([])

        ADC_raw = (data >> 8) & 0xFFFFFF

        if self.adc_units:
            y = ADC_raw.astype(np.float64)
        else:
            ADC_centered = ADC_raw.astype(np.int32) - 0x800000
            y = self.fullscale * ADC_centered / 0x800000

        if self.scale is not None:
            y = y * self.scale

        t0 = self._read_samples / self.fs
        t = t0 + np.arange(data.size, dtype=np.float64) / self.fs

        self._read_samples += data.size
        return t, y

    def close(self) -> None:
        if not self._closed:
            try:
                self._file.close()
            finally:
                self._closed = True

    def remaining_samples(self) -> int:
        return max(0, self.total_samples - self._read_samples)


# ----------------------------------------------------------------------
# Optional Numba
# ----------------------------------------------------------------------

HAVE_NUMBA = False
try:
    from numba import njit  # type: ignore

    print("Numba found; JIT acceleration enabled.")
    HAVE_NUMBA = True
except ImportError:  # pragma: no cover - optional dependency
    print("Numba not found; proceeding without JIT acceleration.")

    def njit(*args, **kwargs):  # type: ignore
        # Dummy decorator if numba is not installed
        def deco(f):
            return f

        return deco


# ----------------------------------------------------------------------
# Units handling (dependency-injected, completely separated)
# ----------------------------------------------------------------------

try:
    import pint  # type: ignore

    _HAVE_PINT = True
except ImportError:  # pragma: no cover - optional dependency
    pint = None  # type: ignore
    _HAVE_PINT = False

NumberLike = Union[int, float]
TimeLike = Union[int, float, str, Any]
FreqLike = Union[int, float, str, Any]


class UnitsAdapter(Protocol):
    """
    Abstract interface for unit handling.

    Implementations must convert arbitrary user inputs to *floats*:
      - Frequencies in Hz
      - Times in seconds

    The rest of the library only sees plain floats.
    """

    def to_frequency_hz(self, value: FreqLike) -> float:
        ...

    def to_time_s(self, value: TimeLike) -> float:
        ...


class SIUnitsAdapter:
    """
    Minimal adapter assuming SI floats.

    - float/int => used as-is (Hz or s)
    - str       => parsed via float(value); no units allowed.

    This is the fallback if pint is not available.
    """

    def to_frequency_hz(self, value: FreqLike) -> float:
        if isinstance(value, str):
            return float(value)
        return float(value)

    def to_time_s(self, value: TimeLike) -> float:
        if isinstance(value, str):
            return float(value)
        return float(value)


if _HAVE_PINT:

    class PintUnitsAdapter:
        """
        Units adapter using pint.

        Allows:
            - "2 kHz"
            - "10 ms"
            - pint.Quantity
            - floats (SI)
        """

        def __init__(
            self, ureg: Optional["pint.UnitRegistry"] = None
        ) -> None:  # type: ignore[name-defined]
            self.ureg = ureg or pint.UnitRegistry()  # type: ignore[arg-type]

        def to_frequency_hz(self, value: FreqLike) -> float:
            if isinstance(value, str):
                q = self.ureg.Quantity(value)
                return q.to(self.ureg.hertz).magnitude
            if isinstance(value, self.ureg.Quantity):
                return value.to(self.ureg.hertz).magnitude
            return float(value)

        def to_time_s(self, value: TimeLike) -> float:
            if isinstance(value, str):
                q = self.ureg.Quantity(value)
                return q.to(self.ureg.second).magnitude
            if isinstance(value, self.ureg.Quantity):
                return value.to(self.ureg.second).magnitude
            return float(value)


def get_default_units_adapter() -> UnitsAdapter:
    """
    Default units adapter:
        - PintUnitsAdapter if pint is installed
        - SIUnitsAdapter otherwise
    """
    if _HAVE_PINT:
        return PintUnitsAdapter()  # type: ignore[return-value]
    return SIUnitsAdapter()


# ----------------------------------------------------------------------
# Result object (dict-like, with convenience methods)
# ----------------------------------------------------------------------


@dataclass
class FilterResult:
    stages: List[np.ndarray]
    params: Dict[str, Any]

    def __getitem__(self, key: str) -> Any:
        if key == "stages":
            return self.stages
        if key == "params":
            return self.params
        raise KeyError(f"Unknown key {key!r}; valid keys are 'stages' and 'params'")

    def get_filtered_waveform(self) -> np.ndarray:
        """Return the last stage (final shaped waveform)."""
        return self.stages[-1]


# Type alias for a filter implementation taking x and returning FilterResult
FilterImpl = Callable[[np.ndarray], FilterResult]


# ----------------------------------------------------------------------
# 1) CR–RC² Filter (single-τ, core implementation with Numba loop)
# from https://inspirehep.net/files/b3953dda9b8dff1e5fa260e274279ec1
# ----------------------------------------------------------------------


@njit(cache=True, fastmath=True)
def _crrc2_inner_loop(
    x: np.ndarray, alpha: float, beta_RC: float, beta_CR: float
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Numba-accelerated inner loop for CR–RC².

    Parameters
    ----------
    x : 1D float64 array
    alpha, beta_RC, beta_CR : floats

    Returns
    -------
    (y1, y2, y3) : three 1D float64 arrays
    """
    N = x.shape[0]
    y1 = np.empty(N, dtype=np.float64)
    y2 = np.empty(N, dtype=np.float64)
    y3 = np.empty(N, dtype=np.float64)

    a = alpha
    b_rc = beta_RC
    b_cr = beta_CR

    x0 = x[0]
    y1[0] = x0
    y2[0] = 0.0
    y3[0] = 0.0

    a = alpha
    b_rc = beta_RC
    b_cr = beta_CR

    for n in range(1, N):
        y1m1 = y1[n - 1]
        xn = x[n]
        y1n = a * y1m1 + b_rc * xn
        y1[n] = y1n

        dy1 = y1n - y1m1
        y2m1 = y2[n - 1]
        y2n = a * y2m1 + b_cr * dy1
        y2[n] = y2n

        dy2 = y2n - y2m1
        y3[n] = a * y3[n - 1] + b_cr * dy2

    return y1, y2, y3


def _crrc2_core(
    x: np.ndarray,
    fs_hz: float,
    rise_time_s: float,
    m: float = 1.0,
) -> FilterResult:
    """
    Core CR–RC² implementation (no units, floats only).
    Uses Numba-accelerated inner loop if numba is available.
    """
    x = np.asarray(x, dtype=np.float64)
    if x.ndim != 1:
        raise ValueError("x must be a 1D array")

    Ts = 1.0 / fs_hz  # [s]
    tau = m * rise_time_s  # [s]

    Ts_over_tau = Ts / tau
    tau_over_Ts = tau / Ts

    alpha = float(np.exp(-Ts_over_tau))
    beta_RC = 1.0 - alpha
    beta_CR = (1.0 - alpha) * tau_over_Ts

    y1, y2, y3 = _crrc2_inner_loop(x, alpha, beta_RC, beta_CR)

    params = {
        "samp_freq_hz": fs_hz,
        "Ts_s": Ts,
        "rise_time_s": rise_time_s,
        "tau_s": tau,
        "m": m,
        "alpha": alpha,
        "beta_RC": beta_RC,
        "beta_CR": beta_CR,
        "tau_over_Ts": tau_over_Ts,
    }
    stages = [y1, y2, y3]
    return FilterResult(stages=stages, params=params)


def crrc2_filter(
    x: Optional[np.ndarray] = None,
    *,
    samp_freq: FreqLike,
    rise_time: TimeLike,
    m: float = 1.0,
    units: Optional[UnitsAdapter] = None,
) -> Union[FilterResult, FilterImpl]:
    """
    User-facing CR–RC² filter.

    Usage
    -----
    1) Immediate filtering:
        result = crrc2_filter(x, samp_freq="2 kHz", rise_time="10 ms")

    2) Preconfigured filter (for later use / triggering):
        filt = crrc2_filter(samp_freq=2000, rise_time=0.01)
        result = filt(x)

    Returns
    -------
    - If x is not None: FilterResult
    - If x is None:     Callable[[np.ndarray], FilterResult]
    """
    units = units or get_default_units_adapter()
    fs_hz = units.to_frequency_hz(samp_freq)
    rt_s = units.to_time_s(rise_time)

    def impl(data: np.ndarray) -> FilterResult:
        return _crrc2_core(data, fs_hz=fs_hz, rise_time_s=rt_s, m=m)

    if x is None:
        return impl
    return impl(np.asarray(x, dtype=np.float64))


# ----------------------------------------------------------------------
# 2) Triangular filter
# ----------------------------------------------------------------------


def _triangular_core(
    x: np.ndarray,
    fs_hz: float,
    rise_time_s: float,
    normalize: bool = True,
) -> FilterResult:
    """
    Symmetric triangular shaping using cumulative sums, float-only core.
    (Already highly optimized via NumPy; loops are minimal.)
    """
    x = np.asarray(x, dtype=np.float64)
    if x.ndim != 1:
        raise ValueError("x must be a 1D array")

    Ts = 1.0 / fs_hz
    L = max(1, int(round(rise_time_s / Ts)))

    N = x.size
    # cumulative sum with S[0] = 0
    S = np.empty(N + 1, dtype=np.float64)
    S[0] = 0.0
    np.cumsum(x, out=S[1:])

    y = np.zeros(N, dtype=np.float64)
    if 2 * L <= N:
        # corresponds to n indices from 2L-1 to N-1
        part = S[2 * L :] - 2.0 * S[L:-L] + S[:-2 * L]
        if normalize:
            part /= float(L)
        y[2 * L - 1 :] = part

    params = {
        "samp_freq_hz": fs_hz,
        "Ts_s": Ts,
        "rise_time_s": rise_time_s,
        "L": L,
        "normalized": normalize,
    }
    stages = [y]
    return FilterResult(stages=stages, params=params)


def triangular_filter(
    x: Optional[np.ndarray] = None,
    *,
    samp_freq: FreqLike,
    rise_time: TimeLike,
    normalize: bool = True,
    units: Optional[UnitsAdapter] = None,
) -> Union[FilterResult, FilterImpl]:
    """
    Triangular shaper (moving-slope) using cumulative sums.

    Same calling convention as crrc2_filter.
    """
    units = units or get_default_units_adapter()
    fs_hz = units.to_frequency_hz(samp_freq)
    rt_s = units.to_time_s(rise_time)

    def impl(data: np.ndarray) -> FilterResult:
        return _triangular_core(data, fs_hz=fs_hz, rise_time_s=rt_s, normalize=normalize)

    if x is None:
        return impl
    return impl(np.asarray(x, dtype=np.float64))


# ----------------------------------------------------------------------
# 3) Trapezoidal filter (Jordanov–Knoll)
# ----------------------------------------------------------------------


@njit(cache=True, fastmath=True)
def _trapezoid_knoll_inner(
    x: np.ndarray,
    k: int,
    l: int,
    M: float,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Numba-accelerated inner loop for trapezoidal shaper.
    """
    N = x.shape[0]
    d = np.zeros(N, dtype=np.float64)
    p = np.zeros(N, dtype=np.float64)
    r = np.zeros(N, dtype=np.float64)
    s = np.zeros(N, dtype=np.float64)

    k_local = k
    l_local = l
    M_local = M

    for n in range(N):
        x0 = x[0]
        xn = x[n]
        xk = x[n - k_local] if n >= k_local else x0
        xl = x[n - l_local] if n >= l_local else x0
        xkl = x[n - k_local - l_local] if n >= k_local + l_local else x0

        dn = xn - xk - xl + xkl
        d[n] = dn

        if n > 0:
            pn = p[n - 1] + dn
        else:
            pn = dn
        p[n] = pn

        rn = pn + M_local * dn
        r[n] = rn

        if n > 0:
            sn = s[n - 1] + rn
        else:
            sn = rn
        s[n] = sn

    return d, p, r, s


def _trapezoid_knoll_core(
    x: np.ndarray,
    fs_hz: float,
    rise_time_s: float,
    flat_top_s: float = 0.0,
    tau_s: Optional[float] = None,
) -> FilterResult:
    """
    Jordanov–Knoll trapezoidal shaper core (float-only).
    """
    x = np.asarray(x, dtype=np.float64)
    if x.ndim != 1:
        raise ValueError("x must be a 1D array")

    Ts = 1.0 / fs_hz

    total_s = rise_time_s + flat_top_s
    k = int(round(rise_time_s / Ts))
    l = int(round(total_s / Ts))

    if k <= 0:
        k = 1
    if l <= 0:
        l = k

    if tau_s is None:
        M = 0.0
    else:
        ratio = tau_s / Ts
        M = ratio - 0.5

    d, p, r, s = _trapezoid_knoll_inner(x, k, l, M)

    params = {
        "samp_freq_hz": fs_hz,
        "Ts_s": Ts,
        "rise_time_s": rise_time_s,
        "flat_top_s": flat_top_s,
        "tau_s": tau_s,
        "k": k,
        "l": l,
        "M": M,
    }
    stages = [d, p, r, s]
    return FilterResult(stages=stages, params=params)


def trapezoid_knoll_filter(
    x: Optional[np.ndarray] = None,
    *,
    samp_freq: FreqLike,
    rise_time: TimeLike,
    flat_top: TimeLike = 0.0,
    tau: Optional[TimeLike] = None,
    units: Optional[UnitsAdapter] = None,
) -> Union[FilterResult, FilterImpl]:
    """
    Jordanov–Knoll trapezoidal shaper.

    Same dual interface as other filters.
    """
    units = units or get_default_units_adapter()
    fs_hz = units.to_frequency_hz(samp_freq)
    rise_s = units.to_time_s(rise_time)
    flat_s = units.to_time_s(flat_top) if flat_top is not None else 0.0
    tau_s: Optional[float] = None
    if tau is not None:
        tau_s = units.to_time_s(tau)

    def impl(data: np.ndarray) -> FilterResult:
        return _trapezoid_knoll_core(
            data,
            fs_hz=fs_hz,
            rise_time_s=rise_s,
            flat_top_s=flat_s,
            tau_s=tau_s,
        )

    if x is None:
        return impl
    return impl(np.asarray(x, dtype=np.float64))


__all__: List[str] = [
    "UnitsAdapter",
    "SIUnitsAdapter",
    "get_default_units_adapter",
    "FilterResult",
    "crrc2_filter",
    "triangular_filter",
    "trapezoid_knoll_filter",
]


if _HAVE_PINT:
    __all__.append("PintUnitsAdapter")


# ----------------------------------------------------------------------
# 4) Baseline RMS estimators (MAD and P²)
# ----------------------------------------------------------------------


@njit(cache=True, fastmath=True)
def _mad_first_diff_segment(dseg: np.ndarray) -> float:
    """
    Numba-accelerated MAD estimator on a first-difference segment.
    Returns sigma_d (the std of the first difference).
    """
    N = dseg.shape[0]
    if N <= 2:
        return 0.0

    # compute |d|
    absd = np.empty(N, dtype=np.float64)
    for i in range(N):
        v = dseg[i]
        if v < 0.0:
            v = -v
        absd[i] = v

    # median(|d|)
    absd_sorted = np.sort(absd)
    if N % 2 == 1:
        m = absd_sorted[N // 2]
    else:
        m = 0.5 * (absd_sorted[N // 2 - 1] + absd_sorted[N // 2])

    return 1.4826 * m  # MAD → sigma_d


# ----------------------------------------------------------------------
# P² median core (Numba-compatible) and RMS computation
# ----------------------------------------------------------------------


@njit(cache=True, fastmath=True)
def _p2_update(
    q: np.ndarray,
    n: np.ndarray,
    np_des: np.ndarray,
    dn: np.ndarray,
    x: float,
):
    """
    Single P² median update step for new sample x = |d[n]|.
    """
    # Step 1: classification
    if x < q[0]:
        q[0] = x
        k = 0
    elif x < q[1]:
        k = 0
    elif x < q[2]:
        k = 1
    elif x < q[3]:
        k = 2
    elif x <= q[4]:
        k = 3
    else:
        q[4] = x
        k = 3

    # Step 2: increment positions
    for i in range(k + 1, 5):
        n[i] += 1
    for i in range(5):
        np_des[i] += dn[i]

    # Step 3: adjust internal markers 1..3
    for i in range(1, 4):
        d = np_des[i] - n[i]
        if (d >= 1.0 and (n[i + 1] - n[i]) > 1) or (d <= -1.0 and (n[i] - n[i - 1]) > 1):
            ds = 1 if d > 0.0 else -1

            # Parabolic estimate
            n_im1 = n[i - 1]
            n_i = n[i]
            n_ip1 = n[i + 1]

            denom = float(n_ip1 - n_im1)
            if denom == 0.0:
                continue

            q_im1 = q[i - 1]
            q_i = q[i]
            q_ip1 = q[i + 1]

            term1_d = float(n_ip1 - n_i)
            term2_d = float(n_i - n_im1)
            if term1_d == 0.0 or term2_d == 0.0:
                continue

            qip = q_i + ds / denom * (
                (n_i - n_im1 + ds) * (q_ip1 - q_i) / term1_d
                + (n_ip1 - n_i - ds) * (q_i - q_im1) / term2_d
            )

            # guardrails: fallback to linear
            if qip <= q_im1 or qip >= q_ip1:
                idx = i + ds
                num = q[idx] - q_i
                den = float(n[idx] - n_i)
                if den != 0.0:
                    qip = q_i + ds * (num / den)
                else:
                    qip = q_i

            q[i] = qip
            n[i] += ds


@njit(cache=True, fastmath=True)
def _baseline_rms_p2_core(y: np.ndarray) -> np.ndarray:
    """
    Numba-accelerated streaming P² median–MAD RMS estimator.

    sigma_x = 1.4826 * median(|d|) / sqrt(2)
    """
    N = y.size
    rms = np.zeros(N, dtype=np.float64)
    if N == 0:
        return rms

    rms[0] = 0.0
    prev = y[0]

    # P² data
    q = np.zeros(5, dtype=np.float64)
    n = np.zeros(5, dtype=np.int64)
    np_des = np.zeros(5, dtype=np.float64)
    dn = np.array([0.0, 0.5, 1.0, 0.5, 0.0], dtype=np.float64)

    buf = np.zeros(5, dtype=np.float64)
    count = 0
    initialized = False

    sqrt2 = np.sqrt(2.0)

    for i in range(1, N):
        # |d|
        dval = y[i] - prev
        if dval < 0.0:
            dval = -dval
        prev = y[i]

        # Not enough samples yet → fill buffer and compute median directly
        if not initialized:
            buf[count] = dval
            count += 1

            if count < 5:
                # direct median over 1,2,3,4 samples
                tmp = np.empty(count, dtype=np.float64)
                for k in range(count):
                    tmp[k] = buf[k]
                tmp_sorted = np.sort(tmp)
                if count % 2 == 1:
                    m = tmp_sorted[count // 2]
                else:
                    m = 0.5 * (tmp_sorted[count // 2 - 1] + tmp_sorted[count // 2])
                sigma_d = 1.4826 * m
                rms[i] = sigma_d / sqrt2
                continue

            # count == 5 → initialize P² markers
            tmp = np.empty(5, dtype=np.float64)
            for k in range(5):
                tmp[k] = buf[k]
            tmp_sorted = np.sort(tmp)

            for k in range(5):
                q[k] = tmp_sorted[k]
                n[k] = k + 1
                np_des[k] = k + 1

            m = q[2]
            sigma_d = 1.4826 * m
            rms[i] = sigma_d / sqrt2
            initialized = True
            continue

        # Fully initialized → apply P² update
        _p2_update(q, n, np_des, dn, dval)
        m = q[2]
        sigma_d = 1.4826 * m
        rms[i] = sigma_d / sqrt2

    return rms


# ----------------------------------------------------------------------
# 4a) RMS Estimator – Solid MAD sliding-window (sparse sampling)
# ----------------------------------------------------------------------


def baseline_rms_mad(
    x: Union[np.ndarray, FilterResult, None] = None,
    *,
    time: TimeLike,
    samp_freq: Optional[FreqLike] = None,
    units: Optional[UnitsAdapter] = None,
    step: Optional[TimeLike] = "200 ms",
) -> Union[Dict[str, np.ndarray], Callable[[Union[np.ndarray, FilterResult]], Dict[str, np.ndarray]]]:
    """
    Running baseline RMS estimator using sparse-window MAD.

    Returns
    -------
    dict with one key:
        "rms" : array of same length as waveform
    """
    units = units or get_default_units_adapter()

    def impl(data: Union[np.ndarray, FilterResult]) -> Dict[str, np.ndarray]:
        # extract waveform and sampling freq
        if isinstance(data, FilterResult):
            y = data.get_filtered_waveform()
            fs_hz = data["params"]["samp_freq_hz"]
        else:
            y = np.asarray(data, dtype=np.float64)
            if samp_freq is None:
                raise ValueError("samp_freq must be provided for ndarray input")
            fs_hz = units.to_frequency_hz(samp_freq)

        N = y.size
        if N == 0:
            return {"rms": np.zeros_like(y)}

        # window & step sizes
        win_s = units.to_time_s(time)
        win_n = max(4, int(round(fs_hz * win_s)))

        step_s = units.to_time_s(step)
        step_n = max(1, int(round(fs_hz * step_s)))

        half = win_n // 2

        d = np.diff(y)
        Nd = d.size

        # anchor points
        anchor_idx = np.arange(0, N, step_n)
        anchor_val = np.zeros(anchor_idx.size, dtype=np.float64)

        for j in range(anchor_idx.size):
            i = anchor_idx[j]
            start = i - half
            if start < 0:
                start = 0
            stop = i + half
            if stop > Nd:
                stop = Nd

            sigma_d = _mad_first_diff_segment(d[start:stop])
            anchor_val[j] = sigma_d / np.sqrt(2.0)

        # interpolate to all samples
        rms = np.interp(
            np.arange(N, dtype=np.float64),
            anchor_idx.astype(np.float64),
            anchor_val,
            left=anchor_val[0],
            right=anchor_val[-1],
        )

        return {"rms": rms}

    if x is None:
        return impl
    return impl(x)


# ----------------------------------------------------------------------
# 4b) RMS Estimator – Streaming P² (fully online)
# ----------------------------------------------------------------------


def baseline_rms_p2(
    x: Union[np.ndarray, FilterResult, None] = None,
    *,
    samp_freq: Optional[FreqLike] = None,
    units: Optional[UnitsAdapter] = None,
) -> Union[Dict[str, np.ndarray], Callable[[Union[np.ndarray, FilterResult]], Dict[str, np.ndarray]]]:
    """
    Streaming baseline RMS estimator using P² median–MAD.

    Returns
    -------
    dict with one key:
        "rms" : array of same length as waveform
    """
    units = units or get_default_units_adapter()

    def impl(data: Union[np.ndarray, FilterResult]) -> Dict[str, np.ndarray]:
        if isinstance(data, FilterResult):
            y = data.get_filtered_waveform()
        else:
            y = np.asarray(data, dtype=np.float64)
            if samp_freq is None:
                raise ValueError("samp_freq must be provided for ndarray input")
            # fs_hz unused but kept for API symmetry

        y = np.asarray(y, dtype=np.float64)
        rms = _baseline_rms_p2_core(y)

        return {"rms": rms}

    if x is None:
        return impl
    return impl(x)


__all__.extend(
    [
        "baseline_rms_mad",
        "baseline_rms_p2",
    ]
)

# ----------------------------------------------------------------------
# 4c) RMS Estimator – Streaming t-Digest median–MAD
# ----------------------------------------------------------------------

try:
    from tdigest import TDigest  # type: ignore

    _HAVE_TDIGEST = True
except ImportError:  # pragma: no cover - optional dependency
    TDigest = None  # type: ignore
    _HAVE_TDIGEST = False


def baseline_rms_tdigest(
    x: Union[np.ndarray, FilterResult, None] = None,
    *,
    samp_freq: Optional[FreqLike] = None,
    units: Optional[UnitsAdapter] = None,
) -> Union[Dict[str, np.ndarray], Callable[[Union[np.ndarray, FilterResult]], Dict[str, np.ndarray]]]:
    """
    Streaming baseline RMS estimator using a t-Digest quantile tracker.

    Converts median(|d|) into RMS via:
        sigma_x = 1.4826 * median(|d|) / sqrt(2)

    Notes
    -----
    - Requires `tdigest` package.
    - Fully streaming, no finite window required.
    - More accurate than P² for non-Gaussian / long-tail noise.
    """
    if not _HAVE_TDIGEST:
        raise ImportError(
            "t-Digest not available. Install it via `pip install tdigest`."
        )

    units = units or get_default_units_adapter()

    def impl(data: Union[np.ndarray, FilterResult]) -> Dict[str, np.ndarray]:
        # Extract waveform
        if isinstance(data, FilterResult):
            y = data.get_filtered_waveform()
        else:
            y = np.asarray(data, dtype=np.float64)
            if samp_freq is None:
                raise ValueError("samp_freq must be provided for ndarray input")
            # fs unused for t-digest (kept for API symmetry)

        y = np.asarray(y, dtype=np.float64)
        N = y.size

        if N == 0:
            return {"rms": np.zeros_like(y)}

        # t-Digest instance
        td = TDigest()

        rms = np.zeros(N, dtype=np.float64)
        prev = y[0]
        rms[0] = 0.0

        sqrt2 = np.sqrt(2.0)

        for i in range(1, N):
            # |d| = |y[n] - y[n-1]|
            dval = y[i] - prev
            prev = y[i]
            if dval < 0.0:
                dval = -dval

            td.update(float(dval))

            # median(|d|)
            m = td.percentile(50)
            # MAD→RMS
            sigma_d = 1.4826 * m
            rms[i] = sigma_d / sqrt2

        return {"rms": rms}

    if x is None:
        return impl
    return impl(x)


__all__.append("baseline_rms_tdigest")


# ----------------------------------------------------------------------
# 5) Triggering on filtered waveforms using a running RMS (batch)
# ----------------------------------------------------------------------


@dataclass
class TriggerResult:
    positions: np.ndarray  # trigger sample index
    thr_to_zc_time: np.ndarray  # time difference [s]
    params: Dict[str, Any]  # everything else (thresholds, etc.)


def trigger(
    x: Union[np.ndarray, FilterResult],
    *,
    threshold_rms: float,
    filter: FilterImpl,
    running_rms_estimate: Callable[[Union[np.ndarray, FilterResult]], Dict[str, np.ndarray]],
    direction: str = "pos2neg",  # "pos2neg" or "neg2pos"
    arm_counts: int = 1,  # must exceed threshold for this many samples
    hold_off: TimeLike = "0 ms",
    units: Optional[UnitsAdapter] = None,
) -> TriggerResult:
    """
    Trigger on filtered waveform w/ running RMS.

    Parameters
    ----------
    x : ndarray or FilterResult
        Raw waveform OR already-filtered result.
    threshold_rms : float
        Multiplier N ⇒ trigger if filtered > N * running RMS.
    filter : FilterImpl
        Preconfigured filter, e.g., crrc2_filter(...).
    running_rms_estimate : Callable
        RMS estimator, e.g., baseline_rms_p2() or baseline_rms_mad(...).
    direction : {"pos2neg", "neg2pos"}
        Sign of zero crossing to mark trigger time.
    arm_counts : int
        Require this many consecutive over-threshold samples before arming trigger.
    hold_off : time interval w/ units
        Disabled region after a trigger.
    """

    # --- resolve units
    units = units or get_default_units_adapter()

    # --- ensure waveform + filter result objects
    if isinstance(x, FilterResult):
        filt_res = x
    else:
        filt_res = filter(np.asarray(x, dtype=np.float64))

    y = filt_res.get_filtered_waveform()
    fs = float(filt_res["params"]["samp_freq_hz"])
    Ts = 1.0 / fs
    N = y.size

    # --- running RMS
    rms_dict = running_rms_estimate(filt_res)
    rms = np.asarray(rms_dict["rms"], dtype=np.float64)
    thr = threshold_rms * rms

    # --- direction selector
    # detect sign change on y: positive->negative or negative->positive
    if direction == "pos2neg":
        def zero_cross(i: int) -> bool:
            return y[i - 1] > 0.0 and y[i] <= 0.0
    elif direction == "neg2pos":
        def zero_cross(i: int) -> bool:
            return y[i - 1] < 0.0 and y[i] >= 0.0
    else:
        raise ValueError("direction must be 'pos2neg' or 'neg2pos'")

    # --- hold-off in samples
    hold_s = units.to_time_s(hold_off)
    hold_n = max(0, int(round(fs * hold_s)))

    trig_positions: List[int] = []
    thr_to_zc: List[float] = []

    # --- state machine
    armed = False
    over_count = 0
    idx = 1  # start from 1 to allow zero-cross check
    last_trigger_idx = -np.inf

    while idx < N:
        # --- if in hold-off
        if idx < last_trigger_idx + hold_n:
            idx += 1
            continue

        # --- if not armed: watch threshold
        if not armed:
            if y[idx] > thr[idx]:
                over_count += 1
                if over_count >= arm_counts:
                    armed = True
                    arm_idx = idx
            else:
                over_count = 0

        else:
            # --- armed -> wait for zero crossing
            if zero_cross(idx):
                trig_positions.append(idx)
                # time from threshold cross (arm) to zero crossing
                thr_to_zc.append((idx - arm_idx) * Ts)
                last_trigger_idx = idx
                armed = False
                over_count = 0

        idx += 1

    return TriggerResult(
        positions=np.asarray(trig_positions, dtype=np.int64),
        thr_to_zc_time=np.asarray(thr_to_zc, dtype=np.float64),
        params={
            "threshold_rms": threshold_rms,
            "direction": direction,
            "arm_counts": arm_counts,
            "hold_off_s": hold_s,
            "fs_hz": fs,
        },
    )


# ----------------------------------------------------------------------
# 6) STREAMING versions of filters, RMS estimators, and trigger
# ----------------------------------------------------------------------

__all__.extend(
    [
        "CupidWaveformReader",
    ]
)


# ---------------------------- Streaming filters -----------------------


class CRRC2FilterStream:
    """
    Streaming CR–RC² filter.

    Maintains state between blocks so that processing chunk-by-chunk gives
    the same recurrence as the batch implementation (up to FP round-off).

    Parameters
    ----------
    fs_hz : float
        Sampling frequency [Hz].
    rise_time_s : float
        CR–RC² rise time [s].
    m : float
        Shape parameter as in batch version.
    """

    def __init__(self, fs_hz: float, rise_time_s: float, m: float = 1.0) -> None:
        fs_hz = float(fs_hz)
        rise_time_s = float(rise_time_s)
        self.fs_hz = fs_hz
        self.rise_time_s = rise_time_s
        Ts = 1.0 / fs_hz
        tau = m * rise_time_s
        Ts_over_tau = Ts / tau
        tau_over_Ts = tau / Ts

        self.alpha = float(np.exp(-Ts_over_tau))
        self.beta_RC = 1.0 - self.alpha
        self.beta_CR = (1.0 - self.alpha) * tau_over_Ts

        self._initialized = False
        self._y1_prev = 0.0
        self._y2_prev = 0.0
        self._y3_prev = 0.0

    def process(self, x_block: np.ndarray) -> np.ndarray:
        """
        Process a block of samples and return the final CR–RC² stage.

        Parameters
        ----------
        x_block : ndarray
            1D waveform block.

        Returns
        -------
        y_block : ndarray
            Filtered waveform for the block (final stage y3).
        """
        x = np.asarray(x_block, dtype=np.float64)
        if x.ndim != 1 or x.size == 0:
            return np.zeros_like(x, dtype=np.float64)

        y1 = np.empty_like(x)
        y2 = np.empty_like(x)
        y3 = np.empty_like(x)

        a = self.alpha
        b_rc = self.beta_RC
        b_cr = self.beta_CR

        if not self._initialized:
            # First sample of first block: match batch core
            x0 = x[0]
            y1_prev = x0
            y2_prev = 0.0
            y3_prev = 0.0
            y1[0] = y1_prev
            y2[0] = y2_prev
            y3[0] = y3_prev
            start = 1
            self._initialized = True
        else:
            # Continue from previous state
            y1_prev = self._y1_prev
            y2_prev = self._y2_prev
            y3_prev = self._y3_prev

            xn = x[0]
            y1_curr = a * y1_prev + b_rc * xn
            dy1 = y1_curr - y1_prev
            y2_curr = a * y2_prev + b_cr * dy1
            dy2 = y2_curr - y2_prev
            y3_curr = a * y3_prev + b_cr * dy2

            y1[0] = y1_curr
            y2[0] = y2_curr
            y3[0] = y3_curr

            y1_prev = y1_curr
            y2_prev = y2_curr
            y3_prev = y3_curr
            start = 1

        for n in range(start, x.size):
            xn = x[n]
            y1_curr = a * y1_prev + b_rc * xn
            dy1 = y1_curr - y1_prev
            y2_curr = a * y2_prev + b_cr * dy1
            dy2 = y2_curr - y2_prev
            y3_curr = a * y3_prev + b_cr * dy2

            y1[n] = y1_curr
            y2[n] = y2_curr
            y3[n] = y3_curr

            y1_prev = y1_curr
            y2_prev = y2_curr
            y3_prev = y3_curr

        self._y1_prev = y1_prev
        self._y2_prev = y2_prev
        self._y3_prev = y3_prev

        return y3


class TriangularFilterStream:
    """
    Streaming triangular filter that is mathematically equivalent to
    the batch _triangular_core implementation.

    Batch formula (for n >= 2L-1):
        y[n] = (S[n+1] - 2*S[n+1-L] + S[n+1-2L]) / L

    where S[k] = sum_{i=0}^{k-1} x[i] (cumulative sum).
    This class keeps a circular buffer of recent S values so it can
    reproduce this exactly, chunk by chunk.
    """

    def __init__(self, fs_hz: float, rise_time_s: float, normalize: bool = True) -> None:
        fs_hz = float(fs_hz)
        rise_time_s = float(rise_time_s)

        self.fs_hz = fs_hz
        self.rise_time_s = rise_time_s
        self.normalize = normalize

        Ts = 1.0 / fs_hz
        L = max(1, int(round(rise_time_s / Ts)))
        self.L = L

        # We need S[n+1], S[n+1-L], S[n+1-2L]
        # Keep last (2L+1) cumulative sums in a ring buffer.
        self._S_buf_size = 2 * L + 1
        self._S_buf = np.zeros(self._S_buf_size, dtype=np.float64)

        self._cumsum = 0.0    # running S[n+1]
        self._n = -1          # last processed sample index (global, start before 0)

    def process(self, x_block: np.ndarray) -> np.ndarray:
        x = np.asarray(x_block, dtype=np.float64)
        if x.ndim != 1 or x.size == 0:
            return np.zeros_like(x, dtype=np.float64)

        y = np.zeros_like(x, dtype=np.float64)
        L = self.L
        Sbuf = self._S_buf
        Ssize = self._S_buf_size

        for i in range(x.size):
            xn = x[i]
            # advance global index
            self._n += 1
            n = self._n

            # update cumulative sum S[n+1]
            self._cumsum += xn
            S_n1 = self._cumsum

            # store S[n+1] in circular buffer
            idx_Sn1 = n % Ssize
            Sbuf[idx_Sn1] = S_n1

            if n >= 2 * L - 1:
                # indices for S[n+1-L] and S[n+1-2L]
                idx_Sn1_L  = (n - L) % Ssize
                idx_Sn1_2L = (n - 2 * L) % Ssize

                S_n1_L  = Sbuf[idx_Sn1_L]
                S_n1_2L = Sbuf[idx_Sn1_2L]

                val = S_n1 - 2.0 * S_n1_L + S_n1_2L
                if self.normalize:
                    val /= float(L)
                y[i] = val
            else:
                # For n < 2L-1, batch triangular filter outputs zero
                y[i] = 0.0

        return y

class TrapezoidKnollFilterStream:
    """
    Streaming Jordanov–Knoll trapezoidal shaper.

    Keeps internal delayed samples to maintain continuity across blocks.
    """

    def __init__(
        self,
        fs_hz: float,
        rise_time_s: float,
        flat_top_s: float = 0.0,
        tau_s: Optional[float] = None,
    ) -> None:
        fs_hz = float(fs_hz)
        rise_time_s = float(rise_time_s)
        flat_top_s = float(flat_top_s)

        self.fs_hz = fs_hz
        Ts = 1.0 / fs_hz

        total_s = rise_time_s + flat_top_s
        k = int(round(rise_time_s / Ts))
        l = int(round(total_s / Ts))

        if k <= 0:
            k = 1
        if l <= 0:
            l = k

        if tau_s is None:
            M = 0.0
        else:
            ratio = tau_s / Ts
            M = ratio - 0.5

        self.k = k
        self.l = l
        self.M = float(M)

        # delay buffer to emulate x[n-k], x[n-l], x[n-k-l]
        self._buf = np.zeros(k + l, dtype=np.float64)
        self._buf_idx = 0
        self._initialized = False

        self._p_prev = 0.0
        self._s_prev = 0.0

    def process(self, x_block: np.ndarray) -> np.ndarray:
        x = np.asarray(x_block, dtype=np.float64)
        if x.ndim != 1 or x.size == 0:
            return np.zeros_like(x, dtype=np.float64)

        N = x.size
        out = np.zeros(N, dtype=np.float64)

        k = self.k
        l = self.l
        M = self.M
        buf = self._buf
        buf_size = buf.size
        idx = self._buf_idx

        p_prev = self._p_prev
        s_prev = self._s_prev

        for i in range(N):
            xn = x[i]

            # current sample in buffer
            buf[idx] = xn

            # indices for delayed samples
            idx_k = (idx - k) % buf_size
            idx_l = (idx - l) % buf_size
            idx_kl = (idx - k - l) % buf_size

            xk = buf[idx_k]
            xl = buf[idx_l]
            xkl = buf[idx_kl]

            dn = xn - xk - xl + xkl

            pn = p_prev + dn
            rn = pn + M * dn
            sn = s_prev + rn

            out[i] = sn

            p_prev = pn
            s_prev = sn
            idx = (idx + 1) % buf_size

        self._buf_idx = idx
        self._p_prev = p_prev
        self._s_prev = s_prev

        return out


__all__.extend(
    [
        "CRRC2FilterStream",
        "TriangularFilterStream",
        "TrapezoidKnollFilterStream",
    ]
)


# ------------------------- Streaming RMS estimators -------------------


class BaselineRMSMadStream:
    """
    Streaming MAD-based RMS estimator (sparse sampling, causal).

    This is *not* identical to the symmetric-window, interpolated
    offline baseline_rms_mad, but it uses the same MAD(|d|)->RMS idea
    in a streaming / low-latency form.

    Parameters
    ----------
    fs_hz : float
        Sampling frequency [Hz].
    win_s : float
        Window length in seconds for MAD over |d|.
    step_s : float
        How often to recompute MAD; RMS is held between updates.
    """

    def __init__(self, fs_hz: float, win_s: float, step_s: float = 0.2) -> None:
        fs_hz = float(fs_hz)
        win_s = float(win_s)
        step_s = float(step_s)

        self.fs_hz = fs_hz
        self.win_n = max(4, int(round(fs_hz * win_s)))
        self.step_n = max(1, int(round(fs_hz * step_s)))

        self._diff_buf = np.zeros(self.win_n, dtype=np.float64)
        self._diff_count = 0
        self._diff_idx = 0

        self._prev_y: Optional[float] = None
        self._sample_index = 0
        self._last_rms = 0.0

    def process(self, y_block: np.ndarray) -> np.ndarray:
        y = np.asarray(y_block, dtype=np.float64)
        if y.ndim != 1 or y.size == 0:
            return np.zeros_like(y, dtype=np.float64)

        rms_block = np.zeros_like(y, dtype=np.float64)
        sqrt2 = np.sqrt(2.0)

        for i, yi in enumerate(y):
            if self._prev_y is None:
                self._prev_y = yi
                rms_block[i] = self._last_rms
                self._sample_index += 1
                continue

            dval = yi - self._prev_y
            self._prev_y = yi
            if dval < 0.0:
                dval = -dval

            self._diff_buf[self._diff_idx] = dval
            self._diff_idx = (self._diff_idx + 1) % self.win_n
            self._diff_count = min(self._diff_count + 1, self.win_n)

            if self._sample_index % self.step_n == 0 and self._diff_count > 2:
                # gather contiguous slice of diffs
                if self._diff_count == self.win_n:
                    if self._diff_idx == 0:
                        dseg = self._diff_buf
                    else:
                        dseg = np.concatenate(
                            (
                                self._diff_buf[self._diff_idx :],
                                self._diff_buf[: self._diff_idx],
                            )
                        )
                else:
                    dseg = self._diff_buf[: self._diff_count]

                sigma_d = _mad_first_diff_segment(dseg)
                self._last_rms = sigma_d / sqrt2

            rms_block[i] = self._last_rms
            self._sample_index += 1

        return rms_block


class BaselineRMSP2Stream:
    """
    Streaming P² median–MAD baseline RMS estimator that is **exactly**
    equivalent to _baseline_rms_p2_core run on the concatenated signal.

    Usage
    -----
    rms_stream = BaselineRMSP2Stream()
    for block in blocks:
        rms_block = rms_stream.process(block)
        # concatenate rms_block to build full RMS

    If you concatenate all rms_block's in order, you get the same result
    as baseline_rms_p2()(full_waveform)["rms"].
    """

    def __init__(self) -> None:
        # state that mirrors _baseline_rms_p2_core
        self._initialized = False       # corresponds to 'initialized'
        self._prev_y: Optional[float] = None
        self._buf = np.zeros(5, dtype=np.float64)
        self._count = 0

        self._q = np.zeros(5, dtype=np.float64)
        self._n = np.zeros(5, dtype=np.int64)
        self._np_des = np.zeros(5, dtype=np.float64)
        self._dn = np.array([0.0, 0.5, 1.0, 0.5, 0.0], dtype=np.float64)

        self._sqrt2 = np.sqrt(2.0)

        # global index i from _baseline_rms_p2_core
        self._global_index = 0

    def reset(self) -> None:
        """Optional helper to restart the stream."""
        self.__init__()

    def process(self, y_block: np.ndarray) -> np.ndarray:
        """
        Process a block of y samples and return the RMS for this block.

        Concatenating all outputs (in order) over all processed blocks
        reproduces _baseline_rms_p2_core(y_full) exactly.
        """
        y = np.asarray(y_block, dtype=np.float64)
        if y.ndim != 1 or y.size == 0:
            return np.zeros_like(y, dtype=np.float64)

        rms_block = np.zeros_like(y, dtype=np.float64)

        for j in range(y.size):
            yi = y[j]

            # --- match _baseline_rms_p2_core: rms[0] = 0, prev = y[0]
            if self._global_index == 0:
                self._prev_y = yi
                rms_block[j] = 0.0
                self._global_index += 1
                continue

            # From here on we are in the "for i in range(1, N)" part
            # of _baseline_rms_p2_core, but i is self._global_index.
            # ----------------- |d| -----------------
            dval = yi - self._prev_y
            self._prev_y = yi
            if dval < 0.0:
                dval = -dval

            # ----------------- Not yet initialized -----------------
            if not self._initialized:
                # fill buf[count]
                self._buf[self._count] = dval
                self._count += 1

                if self._count < 5:
                    # direct median of buf[:count]
                    tmp = np.empty(self._count, dtype=np.float64)
                    for k in range(self._count):
                        tmp[k] = self._buf[k]
                    tmp_sorted = np.sort(tmp)
                    if self._count % 2 == 1:
                        m = tmp_sorted[self._count // 2]
                    else:
                        m = 0.5 * (
                            tmp_sorted[self._count // 2 - 1]
                            + tmp_sorted[self._count // 2]
                        )
                    sigma_d = 1.4826 * m
                    rms_block[j] = sigma_d / self._sqrt2
                    self._global_index += 1
                    continue

                # self._count == 5 → initialize P² markers
                tmp = np.empty(5, dtype=np.float64)
                for k in range(5):
                    tmp[k] = self._buf[k]
                tmp_sorted = np.sort(tmp)

                for k in range(5):
                    self._q[k] = tmp_sorted[k]
                    self._n[k] = k + 1
                    self._np_des[k] = k + 1

                m = self._q[2]
                sigma_d = 1.4826 * m
                rms_block[j] = sigma_d / self._sqrt2
                self._initialized = True
                self._global_index += 1
                continue

            # ----------------- Fully initialized → P² update -----------------
            _p2_update(self._q, self._n, self._np_des, self._dn, dval)
            m = self._q[2]
            sigma_d = 1.4826 * m
            rms_block[j] = sigma_d / self._sqrt2

            self._global_index += 1

        return rms_block

class BaselineRMSTDigestStream:
    """
    Streaming t-Digest median–MAD baseline RMS estimator.

    Requires tdigest package.
    """

    def __init__(self) -> None:
        if not _HAVE_TDIGEST:
            raise ImportError(
                "t-Digest not available. Install it via `pip install tdigest`."
            )
        self._td = TDigest()
        self._prev_y: Optional[float] = None
        self._sqrt2 = np.sqrt(2.0)

    def process(self, y_block: np.ndarray) -> np.ndarray:
        y = np.asarray(y_block, dtype=np.float64)
        if y.ndim != 1 or y.size == 0:
            return np.zeros_like(y, dtype=np.float64)

        rms_block = np.zeros_like(y, dtype=np.float64)

        for i in range(y.size):
            yi = y[i]
            if self._prev_y is None:
                self._prev_y = yi
                rms_block[i] = 0.0
                continue

            dval = yi - self._prev_y
            self._prev_y = yi
            if dval < 0.0:
                dval = -dval

            self._td.update(float(dval))
            m = self._td.percentile(50)
            sigma_d = 1.4826 * m
            rms_block[i] = sigma_d / self._sqrt2

        return rms_block


__all__.extend(
    [
        "BaselineRMSMadStream",
        "BaselineRMSP2Stream",
        "BaselineRMSTDigestStream",
    ]
)


# ------------------------- Streaming trigger engine -------------------


@dataclass
class TriggerEvent:
    position: int          # global sample index
    thr_to_zc_time: float  # time between arm and zero-crossing [s]


class TriggerStream:
    """
    Streaming trigger state machine.

    This class wires together:
        - a streaming filter object (with .process(block) -> y_block)
        - a streaming RMS estimator (with .process(y_block) -> rms_block)

    and finds triggers on the fly as blocks come in.

    Parameters
    ----------
    fs_hz : float
        Sampling frequency [Hz].
    threshold_rms : float
        Trigger threshold in units of RMS (N * sigma).
    filter_stream : object
        Must implement .process(x_block) -> y_block (filtered waveform).
    rms_stream : object
        Must implement .process(y_block) -> rms_block (same length).
    direction : {"pos2neg", "neg2pos"}
        Sign of zero crossing to mark trigger time.
    arm_counts : int
        Require this many consecutive over-threshold samples before arming trigger.
    hold_off_s : float
        Disabled region after a trigger [s].
    """

    def __init__(
        self,
        fs_hz: float,
        threshold_rms: float,
        filter_stream: Any,
        rms_stream: Any,
        direction: str = "pos2neg",
        arm_counts: int = 1,
        hold_off_s: float = 0.0,
    ) -> None:
        fs_hz = float(fs_hz)
        self.fs_hz = fs_hz
        self.Ts = 1.0 / fs_hz

        self.threshold_rms = float(threshold_rms)
        self.filter_stream = filter_stream
        self.rms_stream = rms_stream

        if direction not in ("pos2neg", "neg2pos"):
            raise ValueError("direction must be 'pos2neg' or 'neg2pos'")
        self.direction = direction

        self.arm_counts = int(arm_counts)
        self.hold_n = max(0, int(round(fs_hz * float(hold_off_s))))

        # state
        self._global_idx = 0
        self._last_trigger_idx = -np.inf
        self._armed = False
        self._over_count = 0
        self._arm_idx = 0
        self._prev_y: Optional[float] = None

    def process(self, x_block: np.ndarray) -> tuple[np.ndarray, np.ndarray, List[TriggerEvent]]:
        """
        Process a block of raw samples.

        Returns
        -------
        y_block : ndarray
            Filtered waveform for this block.
        rms_block : ndarray
            Running RMS for this block.
        triggers : list[TriggerEvent]
            Triggers whose positions fall inside this block.
        """
        x = np.asarray(x_block, dtype=np.float64)
        if x.ndim != 1 or x.size == 0:
            return (
                np.zeros_like(x, dtype=np.float64),
                np.zeros_like(x, dtype=np.float64),
                [],
            )

        # Step 1: filter
        y_block = self.filter_stream.process(x)

        # Step 2: running RMS on filtered waveform
        rms_block = self.rms_stream.process(y_block)
        thr_block = self.threshold_rms * rms_block

        triggers: List[TriggerEvent] = []

        for i in range(x.size):
            idx = self._global_idx
            yi = y_block[i]
            thr_i = thr_block[i]

            # hold-off region
            if idx < self._last_trigger_idx + self.hold_n:
                self._global_idx += 1
                self._prev_y = yi
                continue

            if self._prev_y is None:
                self._prev_y = yi
                self._global_idx += 1
                continue

            # check/maintain armed state
            if not self._armed:
                if yi > thr_i:
                    self._over_count += 1
                    if self._over_count >= self.arm_counts:
                        self._armed = True
                        self._arm_idx = idx
                else:
                    self._over_count = 0
            else:
                # armed: look for zero crossing
                if self.direction == "pos2neg":
                    zc = self._prev_y > 0.0 and yi <= 0.0
                else:
                    zc = self._prev_y < 0.0 and yi >= 0.0

                if zc:
                    dt = (idx - self._arm_idx) * self.Ts
                    triggers.append(TriggerEvent(position=idx, thr_to_zc_time=dt))
                    self._last_trigger_idx = idx
                    self._armed = False
                    self._over_count = 0

            self._prev_y = yi
            self._global_idx += 1

        return y_block, rms_block, triggers


__all__.extend(
    [
        "TriggerEvent",
        "TriggerStream",
    ]
)