"""Tests for the voice pitch estimator reference implementation.

The ground-truth recovery tests reproduce the measured verification of
2026-07-31 (synthetic voiced signals through the full pipeline); the VAD
tests apply the arc's standing rule for detection statistics: Monte
Carlo plus an independently computable special case.
"""

import numpy as np
import pytest
from scipy import signal
from scipy.stats import chi2

from vpe import (VoicePitchEstimator, HampelGate, PowerTracker,
                 cepstral_pitch, cepstral_prominence, design_front_end,
                 frame_power, one_sided_periodogram, power_cepstrum)

FS = 16000.0
NFFT = 2048


def voiced(f0, dur=0.5, fs=FS, n_harm=None, noise=0.01, seed=0):
    """Synthetic voiced signal: harmonics with 1/n amplitudes."""
    rng = np.random.default_rng(seed)
    t = np.arange(int(fs * dur)) / fs
    if n_harm is None:
        n_harm = int(4000 // f0)
    x = sum((1.0 / n) * np.cos(2 * np.pi * n * f0 * t + 0.7 * n)
            for n in range(1, n_harm + 1))
    return x + noise * rng.normal(size=x.size)


class TestFrontEnd:
    def test_band_edges_are_minus_3_db(self):
        sos = design_front_end(FS, (50.0, 400.0))
        w, h = signal.sosfreqz(sos, worN=[50.0, 400.0], fs=FS)
        # Butterworth band edges sit at -3.01 dB by construction
        np.testing.assert_allclose(np.abs(h), 1 / np.sqrt(2), rtol=1e-3)

    def test_passband_and_stopband(self):
        sos = design_front_end(FS, (50.0, 400.0))
        w, h = signal.sosfreqz(sos, worN=[141.0, 5.0, 4000.0], fs=FS)
        gains_db = 20 * np.log10(np.abs(h))
        assert gains_db[0] > -0.5           # geometric centre: flat
        assert gains_db[1] < -30            # a decade below the low edge
        assert gains_db[2] < -30            # a decade above the high edge

    def test_stable(self):
        sos = design_front_end()
        z, p, k = signal.sos2zpk(sos)
        assert np.all(np.abs(p) < 1.0)


class TestPeriodogram:
    def test_parseval(self):
        rng = np.random.default_rng(1)
        frame = rng.normal(size=NFFT)
        win = np.hamming(NFFT)
        p = one_sided_periodogram(frame, win)
        windowed = (frame - frame.mean()) * win
        np.testing.assert_allclose(p.sum(), np.sum(windowed ** 2), rtol=1e-10)

    def test_sine_lands_in_its_bin(self):
        k = 40  # 312.5 Hz at 16 kHz / 2048
        t = np.arange(NFFT)
        x = np.cos(2 * np.pi * k * t / NFFT)
        p = one_sided_periodogram(x)
        assert np.argmax(p) == k
        # all energy in that bin: N * A^2 / 2 with A = 1
        np.testing.assert_allclose(p[k], NFFT / 2, rtol=1e-6)


class TestPowerCepstrum:
    def test_pulse_train_peaks_at_its_period(self):
        period = 100
        x = np.zeros(NFFT)
        x[::period] = 1.0
        p = one_sided_periodogram(x)
        c = power_cepstrum(p, NFFT)
        q = 20 + np.argmax(c[20: NFFT // 2])
        assert q == period

    def test_white_noise_cepstral_variance(self):
        # Independent special case: for rectangular-windowed white noise
        # the interior log-periodogram bins are log(exponential), with
        # variance pi^2/6; the cepstrum averages them down to
        # var(c[q]) ~= pi^2 / (6 * nfft) for interior quefrencies.
        nfft, n_mc = 1024, 300
        rng = np.random.default_rng(7)
        qs = slice(40, 320)
        samples = []
        for _ in range(n_mc):
            p = one_sided_periodogram(rng.normal(size=nfft))
            samples.append(power_cepstrum(p, nfft)[qs])
        var = np.var(np.array(samples), axis=0).mean()
        predicted = np.pi ** 2 / (6 * nfft)
        assert abs(var / predicted - 1) < 0.15


class TestCepstralPitch:
    def test_recovers_known_f0_full_pipeline(self):
        # The measured verification of 2026-07-31, as a permanent test:
        # bandpass front end + Hamming periodogram + cepstrum recovers
        # true f0 across the speaking range.
        sos = design_front_end()
        win = np.hamming(NFFT)
        for f0 in [100, 125, 150, 200, 250, 300]:
            x = signal.sosfilt(sos, voiced(f0))
            p = one_sided_periodogram(x[-NFFT:], win)
            c = power_cepstrum(p, NFFT)
            est, _ = cepstral_pitch(c, FS)
            assert abs(est / f0 - 1) < 0.02, f"f0={f0}: got {est:.1f}"

    def test_recovers_f0_without_front_end(self):
        win = np.hamming(NFFT)
        for f0 in [100, 200, 300]:
            p = one_sided_periodogram(voiced(f0)[-NFFT:], win)
            c = power_cepstrum(p, NFFT)
            est, _ = cepstral_pitch(c, FS)
            assert abs(est / f0 - 1) < 0.02

    def test_quefrency_convention(self):
        # A cepstral peak at index q means f0 = fs / q: the integer peak
        # for a known f0 must land at round(fs / f0).
        f0 = 160.0  # fs / f0 = 100 exactly
        win = np.hamming(NFFT)
        p = one_sided_periodogram(voiced(f0)[-NFFT:], win)
        c = power_cepstrum(p, NFFT)
        est, _ = cepstral_pitch(c, FS, refine=False)
        assert round(FS / est) == round(FS / f0)

    def test_parabolic_refinement_beats_grid(self):
        # 173 Hz sits between quefrency gridpoints (fs/173 = 92.5); the
        # raw grid is off by up to f0^2/(2 fs) ~ 0.9 Hz, refinement
        # should do better.
        f0 = 173.0
        win = np.hamming(NFFT)
        p = one_sided_periodogram(voiced(f0)[-NFFT:], win)
        c = power_cepstrum(p, NFFT)
        est_ref, _ = cepstral_pitch(c, FS, refine=True)
        assert abs(est_ref - f0) < 1.0

    def test_band_incompatible_raises(self):
        with pytest.raises(ValueError):
            cepstral_pitch(np.zeros(8), FS)


class TestProminence:
    def test_voiced_far_above_noise(self):
        win = np.hamming(NFFT)
        p = one_sided_periodogram(voiced(150.0)[-NFFT:], win)
        z_voiced = cepstral_prominence(power_cepstrum(p, NFFT), FS)
        rng = np.random.default_rng(3)
        p_n = one_sided_periodogram(rng.normal(size=NFFT), win)
        z_noise = cepstral_prominence(power_cepstrum(p_n, NFFT), FS)
        assert z_voiced > 3 * z_noise
        assert z_voiced > 10

    def test_noise_prominence_matches_gaussian_max(self):
        # ~280 approximately independent Gaussian samples in the search
        # window: the expected maximum z-score is near sqrt(2 ln 280)
        # ~ 3.4 and should essentially never reach 6.
        rng = np.random.default_rng(11)
        zs = []
        for _ in range(100):
            p = one_sided_periodogram(rng.normal(size=NFFT))
            zs.append(cepstral_prominence(power_cepstrum(p, NFFT), FS))
        zs = np.array(zs)
        assert 2.0 < zs.mean() < 5.0
        assert zs.max() < 6.0


class TestFramePower:
    def test_filtered_noise_effective_dof(self):
        # The VAD sees the frame AFTER the 50-400 Hz front end, so its
        # power statistic has the Welch-Satterthwaite effective DOF
        # 2 * B_eq * T ~= 132, not 2048.  Pin the spread of the deployed
        # statistic so the page's calibration stays honest.
        rng = np.random.default_rng(17)
        sos = design_front_end()
        x = signal.sosfilt(sos, rng.normal(size=200 * NFFT))
        frames = x.reshape(-1, NFFT)
        p = np.mean(frames ** 2, axis=1)
        dof = 2 / (p.std() / p.mean()) ** 2
        assert 100 < dof < 175

    def test_chi_square_tail(self):
        # Standing rule: MC plus independent special case.  For white
        # Gaussian noise, n * mean(x^2) / sigma^2 ~ chi2(n), so the
        # exceedance probability of any fixed threshold is closed-form.
        n, sigma, n_mc = 128, 1.0, 20000
        rng = np.random.default_rng(5)
        frames = rng.normal(0, sigma, size=(n_mc, n))
        powers = np.mean(frames ** 2, axis=1)
        threshold = 1.2  # in power units
        p_mc = np.mean(powers > threshold)
        p_exact = chi2.sf(n * threshold / sigma ** 2, df=n)
        assert abs(p_mc - p_exact) < 0.01
        assert p_exact > 0.02  # threshold chosen so the tail is testable


class TestPowerTracker:
    def test_fast_attack_slow_decay(self):
        tr = PowerTracker(alpha=0.99)
        for _ in range(50):
            tr.update(1.0)
        lo, hi = tr.update(100.0)   # attack: max jumps immediately
        assert hi == 100.0
        lo, hi = tr.update(1.0)     # decay: max relaxes slowly
        assert 95.0 < hi < 100.0
        assert lo <= 1.0

    def test_rejects_bad_alpha(self):
        with pytest.raises(ValueError):
            PowerTracker(alpha=1.5)


class TestHampelGate:
    def test_spike_replaced_by_median(self):
        gate = HampelGate(window=5, k=4.0)
        for v in [100, 101, 99, 100, 101]:
            gate.process(v)
        assert gate.process(200.0) == pytest.approx(100.0, abs=1.5)

    def test_clean_track_passes(self):
        gate = HampelGate(window=5, k=4.0)
        track = [100, 102, 104, 106, 108, 110, 112]
        outputs = [gate.process(v) for v in track]
        np.testing.assert_allclose(outputs, track)

    def test_rejects_tiny_window(self):
        with pytest.raises(ValueError):
            HampelGate(window=2)


class TestVoicePitchEstimator:
    def test_streaming_equals_batch(self):
        # 16384 samples = four whole 4096-sample buffers, so both paths
        # see the identical frame sequence
        x = voiced(150.0, dur=1.024)
        est_a = VoicePitchEstimator()
        est_a.process(x)
        est_b = VoicePitchEstimator()
        for start in range(0, x.size - 4096 + 1, 4096):
            est_b.process(x[start: start + 4096])
        assert est_a.f0 == pytest.approx(est_b.f0, rel=1e-12)
        assert est_a.voiced == est_b.voiced

    def test_tracks_voiced_episode(self):
        rng = np.random.default_rng(9)
        fs = int(FS)
        quiet = 0.01 * rng.normal(size=fs)          # 1 s noise
        speech = voiced(180.0, dur=1.0, seed=2)     # 1 s voiced
        est = VoicePitchEstimator()
        est.process(quiet)
        assert not est.voiced
        est.process(speech)
        assert est.voiced
        assert abs(est.f0 / 180.0 - 1) < 0.02
        # pitch holds (not decays) through the following silence
        est.process(0.01 * rng.normal(size=fs))
        assert not est.voiced
        assert abs(est.f0 / 180.0 - 1) < 0.02

    def test_noise_only_rarely_voiced(self):
        rng = np.random.default_rng(13)
        est = VoicePitchEstimator()
        flags = []
        for _ in range(40):
            est.process(0.01 * rng.normal(size=4096))
            flags.append(est.voiced)
        assert np.mean(flags) < 0.1

    def test_pitch_resets_after_silence_gap(self):
        # A new utterance at a different pitch must not be "corrected"
        # toward the previous utterance's median or blended with its EMA:
        # the first voiced report after a gap should already be the new
        # pitch.
        rng = np.random.default_rng(21)
        est = VoicePitchEstimator()
        est.process(0.01 * rng.normal(size=8192))
        est.process(voiced(180.0, dur=1.024, seed=4))
        assert est.voiced and abs(est.f0 / 180.0 - 1) < 0.02
        est.process(0.01 * rng.normal(size=16384))   # ~1 s silence gap
        assert not est.voiced
        speech2 = voiced(140.0, dur=1.024, seed=5)
        first_voiced_f0 = None
        for start in range(0, speech2.size - 4096 + 1, 4096):
            est.process(speech2[start: start + 4096])
            if est.voiced:
                first_voiced_f0 = est.f0
                break
        assert first_voiced_f0 is not None
        assert abs(first_voiced_f0 / 140.0 - 1) < 0.02

    def test_digital_silence_is_finite_and_unvoiced(self):
        # All-zero input (muted mic, DMA warm-up): no NaN anywhere, and
        # the frame is unvoiced.
        p = one_sided_periodogram(np.zeros(NFFT))
        c = power_cepstrum(p, NFFT)
        assert np.all(np.isfinite(c))
        est = VoicePitchEstimator()
        est.process(np.zeros(16384))
        assert not est.voiced
        assert np.isfinite(est.f0) and np.isfinite(est.prominence)

    def test_hampel_absorbs_octave_glitch(self):
        # A pitch track with a single octave error: the gate holds the
        # median instead of letting the glitch through.
        gate = HampelGate(window=11, k=4.0)
        for _ in range(11):
            gate.process(150.0)
        assert gate.process(300.0) == pytest.approx(150.0)

    def test_rejects_bad_params(self):
        with pytest.raises(ValueError):
            VoicePitchEstimator(nfft=1000)
        with pytest.raises(ValueError):
            VoicePitchEstimator(alpha=1.0)
        with pytest.raises(ValueError):
            VoicePitchEstimator(vad_beta=0.0)