"""Speaking-voice pitch estimator: the estimation & detection arc in one system.

Reference implementation of a classical voice pitch estimator for the
50-400 Hz speaking-pitch band, structured the way the embedded companion
runs on the metal:

    bandpass front end -> averaged one-sided periodogram -> power cepstrum
    -> quefrency peak -> f0, gated by an energy + pitch-band voice activity
    detector and a Hampel gate on the pitch track.

Conventions
-----------
* One-sided periodogram bins carry factor 2 except DC and Nyquist, so that
  ``sum(P)`` equals the frame's total energy (Parseval; tested).
* A cepstral peak at quefrency index q means a fundamental of ``f0 = fs / q``;
  the search window for a pitch band [f_lo, f_hi] is
  ``q in [fs / f_hi, fs / f_lo]``.  Quefrency is time-like: higher pitch,
  shorter period, smaller q.
* Filter coefficients follow the SciPy SOS convention (ADR-005).
"""

from __future__ import annotations

from collections import deque

import numpy as np
from scipy import signal


def design_front_end(fs: float = 16000.0,
                     band: tuple[float, float] = (50.0, 400.0)) -> np.ndarray:
    """Design the bandpass front end: 2nd-order Butterworth, two SOS.

    The band edges are the -3 dB points by Butterworth construction
    (tested, not assumed).  Returns SOS in SciPy convention, directly
    usable by ``scipy.signal.sosfilt`` and pasteable to the embedded
    biquad cascade after factoring out the section-1 scale.
    """
    return signal.butter(N=2, Wn=band, btype="bandpass", output="sos", fs=fs)


def one_sided_periodogram(frame: np.ndarray,
                          window: np.ndarray | None = None) -> np.ndarray:
    """One-sided periodogram of a real frame, energy-preserving.

    P[j] = c_j |X_j|^2 / N with c = 1 at DC and Nyquist, 2 elsewhere,
    so ``sum(P) == sum(frame**2)`` exactly (after windowing, of the
    windowed frame).  The frame mean is removed before windowing, as on
    the metal, so P[0] is essentially zero.
    """
    x = np.asarray(frame, dtype=float)
    n = x.size
    x = x - x.mean()
    if window is not None:
        x = x * window
    spec = np.fft.rfft(x)
    p = np.abs(spec) ** 2 / n
    p[1:] *= 2.0
    if n % 2 == 0:
        p[-1] /= 2.0  # Nyquist bin appears once in the full spectrum
    return p


def power_cepstrum(periodogram: np.ndarray, nfft: int,
                   floor_rel: float = 1e-7) -> np.ndarray:
    """Real cepstrum of the power spectrum: IFFT of the log-periodogram.

    ``np.fft.irfft`` on the (real) one-sided log-periodogram performs the
    even extension implicitly, which is exactly the double-sided
    log-spectrum a real signal has.  Returns the one-sided cepstrum,
    quefrency indices 0..nfft/2, in log-power units.

    Before the log, bins are floored at ``floor_rel`` times the peak bin
    (default -70 dB).  Bins that far down carry window leakage and
    arithmetic noise, not signal structure, and unfloored they dominate
    the log-spectrum with meaningless deep excursions.  With the floor
    in place, the white-noise cepstrum has the clean variance
    pi^2 / (6 * nfft) per interior quefrency (tested), which is what
    makes the voicing statistic below calibratable.
    """
    p = np.asarray(periodogram, dtype=float)
    # absolute tiny under the relative floor: an all-zero frame (muted
    # microphone, DMA warm-up) must yield a finite cepstrum, not NaN
    floor = max(p.max() * floor_rel, np.finfo(float).tiny)
    log_p = np.log(np.maximum(p, floor))
    c = np.fft.irfft(log_p, n=nfft)
    return c[: nfft // 2 + 1]


def cepstral_pitch(cepstrum: np.ndarray, fs: float,
                   band: tuple[float, float] = (50.0, 400.0),
                   refine: bool = True,
                   octave_ratio: float = 0.5) -> tuple[float, float]:
    """Read fundamental frequency off a cepstral peak.

    Searches quefrency indices ``q in [fs / band[1], fs / band[0]]`` and
    returns ``(f0, peak_value)`` with ``f0 = fs / q``.

    Two refinements over the bare argmax, both standard practice in
    cepstral pitch tracking:

    * **Subharmonic disambiguation.**  A voiced cepstrum peaks at the
      pitch quefrency q0 *and* its rahmonics 2 q0, 3 q0, ...; window
      leakage can push a rahmonic a hair above the fundamental, which
      would halve the reported pitch (the classic octave-down error).
      If the half-quefrency neighbourhood of the winning peak holds a
      peak at least ``octave_ratio`` times as tall, that lower quefrency
      wins instead.
    * **Parabolic interpolation** (``refine=True``) sharpens q below the
      integer grid; the raw grid's pitch resolution is f0^2 / fs (about
      5.6 Hz at 300 Hz and 16 kHz).
    """
    q_lo = int(np.ceil(fs / band[1]))
    q_hi = int(np.floor(fs / band[0]))
    q_hi = min(q_hi, cepstrum.size - 2)  # keep q+1 in range for refinement
    if q_lo < 1 or q_lo >= q_hi:
        raise ValueError("pitch band incompatible with cepstrum length")
    seg = cepstrum[q_lo: q_hi + 1]
    q = q_lo + int(np.argmax(seg))

    # Walk down the rahmonic ladder while a comparable peak sits at q/2.
    while True:
        centre = q // 2
        cands = [c for c in (centre - 1, centre, centre + 1) if c >= q_lo]
        if not cands or q // 2 < q_lo:
            break
        q_half = max(cands, key=lambda i: cepstrum[i])
        if cepstrum[q_half] > 0 and \
                cepstrum[q_half] >= octave_ratio * cepstrum[q]:
            q = q_half
        else:
            break

    peak = float(cepstrum[q])
    if refine and 0 < q < cepstrum.size - 1:
        ym, y0, yp = cepstrum[q - 1], cepstrum[q], cepstrum[q + 1]
        denom = ym - 2 * y0 + yp
        if denom < 0:  # genuine maximum; offset in (-1/2, 1/2)
            q = q + 0.5 * (ym - yp) / denom
    return fs / q, peak


def cepstral_prominence(cepstrum: np.ndarray, fs: float,
                        band: tuple[float, float] = (50.0, 400.0)) -> float:
    """Robust z-score of the tallest cepstral peak in the pitch band.

    The voicing evidence carried by the cepstrum: how far the candidate
    peak stands above the surrounding quefrency samples, in units of
    their robust spread (1.4826 * MAD estimates sigma for Gaussian
    data).  For white noise the cepstral samples are approximately
    Gaussian with variance pi^2 / (6 * nfft) (tested against Monte
    Carlo), so the false-voicing rate of a threshold on this statistic
    is calibratable; voiced speech drives it far above any sensible
    threshold.
    """
    q_lo = int(np.ceil(fs / band[1]))
    q_hi = min(int(np.floor(fs / band[0])), cepstrum.size - 1)
    seg = np.asarray(cepstrum[q_lo: q_hi + 1], dtype=float)
    med = np.median(seg)
    mad = np.median(np.abs(seg - med))
    if mad == 0:
        return 0.0
    return float((seg.max() - med) / (1.4826 * mad))


def frame_power(frame: np.ndarray) -> float:
    """Average power of a frame, the VAD's test statistic."""
    x = np.asarray(frame, dtype=float)
    return float(np.mean(x ** 2))


class PowerTracker:
    """Fast-attack, slow-decay trackers for the noise floor and speech peak.

    The minimum tracker follows a drop immediately and relaxes upward
    with time constant ``1 / (1 - alpha)`` updates; the maximum tracker
    mirrors it.  Together they estimate the signal's dynamic range
    without any calibration, which is what lets the VAD threshold adapt
    to the environment (the CFAR idea in miniature; see the
    detection-theory page for the fixed-false-alarm version).
    """

    def __init__(self, alpha: float = 0.9999):
        if not 0.0 < alpha < 1.0:
            raise ValueError("alpha must be in (0, 1)")
        self.alpha = alpha
        self.min_power: float | None = None
        self.max_power: float | None = None

    def update(self, power: float) -> tuple[float, float]:
        a = self.alpha
        if self.min_power is None:
            self.min_power = self.max_power = power
        else:
            if power < self.min_power:
                self.min_power = power
            else:
                self.min_power = a * self.min_power + (1 - a) * power
            if power > self.max_power:
                self.max_power = power
            else:
                self.max_power = a * self.max_power + (1 - a) * power
        return self.min_power, self.max_power


class HampelGate:
    """Median/MAD gate for the pitch track (Hampel identifier).

    Same statistic as the outlier-detection page's ``OutlierDetectorMAD``:
    a sample farther than ``k`` MADs from the window median is replaced
    by that median instead of passing through.  On a pitch track this
    absorbs octave errors and single-frame glitches without lag on
    legitimate pitch movement.
    """

    def __init__(self, window: int = 11, k: float = 4.0):
        if window < 3:
            raise ValueError("window must be at least 3")
        self.buffer: deque[float] = deque(maxlen=window)
        self.k = k

    def process(self, x: float) -> float:
        out = x
        if len(self.buffer) == self.buffer.maxlen:
            buf = np.array(self.buffer)
            med = float(np.median(buf))
            mad = float(np.median(np.abs(buf - med)))
            if abs(x - med) > self.k * mad:
                out = med
        # the raw sample enters the window either way, so a sustained
        # level change works its way into the median within half a window
        self.buffer.append(x)
        return out


class VoicePitchEstimator:
    """Streaming voice pitch estimator, faithful to the embedded pipeline.

    Feed it successive sample buffers (any length >= nfft); it maintains
    filter state, an exponentially averaged periodogram (a running Welch
    estimate), the VAD's power trackers, and the Hampel-gated pitch
    track across calls, exactly as the C implementation does across
    I2S DMA buffers.

    Parameters mirror the hardware defaults: fs = 16 kHz, 2048-point
    Hamming frames, 50% overlap, 50-400 Hz pitch band.
    ``alpha`` is the periodogram/pitch EMA weight on the *old* value
    (alpha = 0 disables smoothing).
    """

    def __init__(self, fs: float = 16000.0, nfft: int = 2048,
                 band: tuple[float, float] = (50.0, 400.0),
                 alpha: float = 0.1, vad_beta: float = 0.5,
                 vad_alpha: float = 0.9999, vad_zmin: float = 5.0,
                 hampel_window: int = 11, hampel_k: float = 4.0,
                 sos: np.ndarray | None = None):
        if nfft & (nfft - 1):
            raise ValueError("nfft must be a power of two")
        if not 0.0 <= alpha < 1.0:
            raise ValueError("alpha must be in [0, 1)")
        if not 0.0 < vad_beta < 1.0:
            raise ValueError("vad_beta must be in (0, 1)")
        self.fs = fs
        self.nfft = nfft
        self.hop = nfft // 2
        self.band = band
        self.alpha = alpha
        self.vad_beta = vad_beta
        self.vad_zmin = vad_zmin
        self.window = np.hamming(nfft)
        self.sos = design_front_end(fs, band) if sos is None else sos
        self._zi = np.zeros((self.sos.shape[0], 2))
        self._tail = np.empty(0)
        self._psd: np.ndarray | None = None
        self.tracker = PowerTracker(vad_alpha)
        self.gate = HampelGate(hampel_window, hampel_k)
        self.voiced = False
        self.f0 = 0.0
        self.magnitude = 0.0
        self.prominence = 0.0
        self._unvoiced_run = 0
        self._restart = True

    def process(self, samples: np.ndarray) -> tuple[bool, float, float]:
        """Consume a buffer; return (voiced, f0 [Hz], cepstral peak).

        The returned f0 is the smoothed, Hampel-gated estimate; it holds
        its last value through unvoiced stretches rather than decaying.
        """
        x = np.asarray(samples, dtype=float)
        filtered, self._zi = signal.sosfilt(self.sos, x, zi=self._zi)
        buf = np.concatenate([self._tail, filtered])

        n_frames = 1 + (buf.size - self.nfft) // self.hop if buf.size >= self.nfft else 0
        for m in range(n_frames):
            frame = buf[m * self.hop: m * self.hop + self.nfft]
            self._update_frame(frame)
        consumed = n_frames * self.hop
        self._tail = buf[consumed:]
        return self.voiced, self.f0, self.magnitude

    def _update_frame(self, frame: np.ndarray) -> None:
        p_inst = one_sided_periodogram(frame, self.window)
        if self._psd is None or self.alpha == 0.0:
            self._psd = p_inst
        else:
            self._psd = self.alpha * self._psd + (1 - self.alpha) * p_inst

        cep = power_cepstrum(self._psd, self.nfft)
        f0_raw, peak = cepstral_pitch(cep, self.fs, self.band)
        self.prominence = cepstral_prominence(cep, self.fs, self.band)

        power = frame_power(frame)
        floor, ceil = self.tracker.update(power)
        eps = np.finfo(float).tiny
        log_p = 10 * np.log10(power + eps)
        log_floor = 10 * np.log10(floor + eps)
        log_ceil = 10 * np.log10(ceil + eps)
        threshold = log_floor + self.vad_beta * (log_ceil - log_floor)
        self.voiced = bool(log_p > threshold
                           and self.prominence > self.vad_zmin)

        if not self.voiced:
            # after a real silence gap the previous utterance's pitch
            # context is stale: forget the gate window so a new utterance
            # at a different pitch is not "corrected" toward the old one.
            # self.f0 itself keeps holding the last display value.
            self._unvoiced_run += 1
            if self._unvoiced_run == 3:
                self.gate.buffer.clear()
                self._restart = True
            return

        self._unvoiced_run = 0
        f0_gated = self.gate.process(f0_raw)
        if self._restart or self.f0 <= 0.0 or self.alpha == 0.0:
            self.f0 = f0_gated
            self._restart = False
        else:
            self.f0 = self.alpha * self.f0 + (1 - self.alpha) * f0_gated
        self.magnitude = peak