code/tests/test_plan_extract.py

# test_plan_extract.py: The prosody plan extractor and quantiser.

from __future__ import annotations

import json
import math
import sys
from pathlib import Path

import numpy as np
import pytest

REPO = Path(__file__).resolve().parent.parent
for p in (str(REPO), str(REPO / "scripts")):
    if p not in sys.path:
        sys.path.insert(0, p)

from plan_extract import (
    DUR_EDGES_FRAMES,
    FRAME_RATE_HZ,
    MIN_WORDS_FOR_SPEAKER_BINS,
    N_ENERGY_BINS,
    N_PITCH_BINS,
    PITCH_UNVOICED_BIN,
    PLAN_WORD_TOKEN,
    PlanAlignmentError,
    PlanExtractError,
    PlanExtractor,
    PlanQuantiseError,
    PlanTokenError,
    dequantise_plan,
    extract_word_plan,
    fit_speaker_stats,
    normalise_transcript,
    plan_frames,
    plan_to_tokens,
    plan_vocab,
    quantise_plan,
    tokens_to_plan,
    verify_config_vocab,
    word_features_from_spans,
)

SR = 16000


def _harmonic(f0_hz: float, dur_s: float, amp: float = 0.3, sr: int = SR) -> np.ndarray:
    t = np.arange(int(round(sr * dur_s))) / sr
    ph = 2 * math.pi * f0_hz * t
    y = np.sin(ph) + 0.5 * np.sin(2 * ph) + 0.25 * np.sin(3 * ph)
    return (amp * y / np.abs(y).max()).astype(np.float64)


def _silence(dur_s: float, sr: int = SR) -> np.ndarray:
    return np.zeros(int(round(sr * dur_s)), dtype=np.float64)


def _spans_from_durations(durations, words=None, conf=1.0):
    out, t = [], 0.0
    for i, d in enumerate(durations):
        name = words[i] if words else f"w{i}"
        out.append((name, t, t + d, conf))
        t += d
    return out


def _synthetic_words(n, *, rng_seed=0, f0_centre=180.0, f0_octaves=0.8,
                     rms_centre=0.05, rms_db_spread=18.0,
                     dur_lo=0.12, dur_hi=0.9, frame_rate=FRAME_RATE_HZ):
    rng = np.random.default_rng(rng_seed)
    words = []
    for i in range(n):
        oct_off = f0_octaves * (rng.random() - 0.5)
        f0 = f0_centre * (2.0 ** oct_off)
        rms = rms_centre * (10.0 ** (rms_db_spread * (rng.random() - 0.5) / 20.0))
        dur_s = float(dur_lo + (dur_hi - dur_lo) * rng.random())
        words.append({
            "word": f"w{i}", "idx": i, "t_start": 0.0, "t_end": dur_s,
            "dur_s": dur_s, "dur_frames": max(1, int(round(dur_s * frame_rate))),
            "f0_mean_hz": float(f0), "voiced_frac": 0.9, "n_voiced_frames": 10,
            "energy_rms": float(rms), "conf": 0.9, "silent": False,
        })
    return words


def test_known_f0_per_word_is_recovered():
    audio = np.concatenate([_harmonic(220.0, 1.0), _harmonic(440.0, 1.0)])
    spans = _spans_from_durations([1.0, 1.0], words=["low", "high"])
    feats = word_features_from_spans(audio, SR, spans)

    assert len(feats) == 2
    lo, hi = feats
    assert lo["voiced_frac"] > 0.8, lo
    assert hi["voiced_frac"] > 0.8, hi
    assert abs(lo["f0_mean_hz"] - 220.0) / 220.0 < 0.03, lo["f0_mean_hz"]
    assert abs(hi["f0_mean_hz"] - 440.0) / 440.0 < 0.03, hi["f0_mean_hz"]


def test_higher_f0_lands_in_a_higher_pitch_bin():
    audio = np.concatenate([_harmonic(220.0, 1.0), _harmonic(440.0, 1.0)])
    spans = _spans_from_durations([1.0, 1.0], words=["low", "high"])
    q = quantise_plan(word_features_from_spans(audio, SR, spans))

    assert q[0]["bin_source"] == "global"
    assert q[1]["pitch_bin"] > q[0]["pitch_bin"], [w["pitch_bin"] for w in q]
    assert all(0 <= w["pitch_bin"] < N_PITCH_BINS for w in q)


@pytest.mark.parametrize("dur_s", [0.10, 0.25, 0.5, 0.93, 1.5])
def test_known_duration_to_frames_within_one_frame(dur_s):
    audio = _harmonic(200.0, dur_s + 0.2)
    feats = word_features_from_spans(audio, SR, [("w", 0.0, dur_s)])
    expected = dur_s * FRAME_RATE_HZ
    assert abs(feats[0]["dur_frames"] - expected) <= 1.0, (
        feats[0]["dur_frames"], expected)


def test_duration_bins_are_monotone_in_duration():
    durs = [0.08, 0.15, 0.3, 0.6, 1.2, 2.0]
    words = [{"word": f"w{i}", "dur_s": d,
              "dur_frames": max(1, int(round(d * FRAME_RATE_HZ))),
              "f0_mean_hz": 180.0, "energy_rms": 0.05}
             for i, d in enumerate(durs)]
    bins = [w["dur_bin"] for w in quantise_plan(words)]
    assert bins == sorted(bins), bins
    assert bins[0] < bins[-1], bins
    assert all(0 <= b < len(DUR_EDGES_FRAMES) + 1 for b in bins)


def test_plan_frames_sums_the_duration_tier():
    words = [{"word": "a", "dur_s": 0.5, "dur_frames": 11,
              "f0_mean_hz": 180.0, "energy_rms": 0.05},
             {"word": "b", "dur_s": 0.25, "dur_frames": 5,
              "f0_mean_hz": 190.0, "energy_rms": 0.05}]
    q = quantise_plan(words)
    assert plan_frames(q) == 16
    from_tokens = plan_frames(plan_to_tokens(q))
    assert 0.6 * 16 <= from_tokens <= 1.6 * 16, from_tokens


def test_duration_tier_reconstructs_the_utterance_length():
    words = _synthetic_words(180, rng_seed=2)
    q = quantise_plan(words)
    true_frames = plan_frames(q)
    est = plan_frames(plan_to_tokens(q))
    assert abs(est - true_frames) / true_frames < 0.10, (est, true_frames)


def test_four_times_amplitude_gives_higher_energy_bin():
    audio = np.concatenate([_harmonic(200.0, 1.0, amp=0.08),
                            _harmonic(200.0, 1.0, amp=0.32)])
    spans = _spans_from_durations([1.0, 1.0], words=["quiet", "loud"])
    feats = word_features_from_spans(audio, SR, spans)

    ratio = feats[1]["energy_rms"] / feats[0]["energy_rms"]
    assert abs(ratio - 4.0) < 0.05, ratio
    q = quantise_plan(feats)
    assert q[1]["energy_bin"] > q[0]["energy_bin"], [w["energy_bin"] for w in q]


def _pearson(a, b):
    a = np.asarray(a, dtype=np.float64)
    b = np.asarray(b, dtype=np.float64)
    return float(np.corrcoef(a, b)[0, 1])


def test_quantiser_roundtrip_preserves_the_contour():
    words = _synthetic_words(180, rng_seed=0)
    stats = fit_speaker_stats([words], speaker_id="spk")
    assert stats["sufficient"], stats["reason"]

    q = quantise_plan(words, stats)
    dq = dequantise_plan(q, stats)

    r_pitch = _pearson([w["pitch_val"] for w in dq], [w["pitch_hat"] for w in dq])
    r_energy = _pearson([w["energy_val"] for w in dq], [w["energy_hat"] for w in dq])
    r_dur = _pearson(np.log([w["dur_frames"] for w in dq]),
                     np.log([w["dur_hat_frames"] for w in dq]))
    assert r_pitch >= 0.8, r_pitch
    assert r_energy >= 0.8, r_energy
    assert r_dur >= 0.8, r_dur


def test_quantile_bins_are_used_uniformly():
    words = _synthetic_words(320, rng_seed=1)
    stats = fit_speaker_stats([words], speaker_id="spk")
    q = quantise_plan(words, stats)
    occ = np.bincount([w["pitch_bin"] for w in q], minlength=N_PITCH_BINS)
    assert (occ[:N_PITCH_BINS] > 0).all(), occ
    assert occ.max() < 0.25 * len(words), occ


def test_within_speaker_binning_equates_relative_prosody_across_registers():
    n = 120
    rel = np.linspace(-0.5, 0.5, n)
    rng = np.random.default_rng(7)
    rms = 0.05 * 10 ** (12.0 * (rng.random(n) - 0.5) / 20.0)

    def mk(base_hz):
        return [{"word": f"w{i}", "dur_s": 0.3,
                 "dur_frames": max(1, int(round(0.3 * FRAME_RATE_HZ))),
                 "f0_mean_hz": float(base_hz * 2.0 ** rel[i]),
                 "energy_rms": float(rms[i])} for i in range(n)]

    low_voice = mk(110.0)
    high_voice = mk(220.0)

    s_low = fit_speaker_stats([low_voice], speaker_id="low")
    s_high = fit_speaker_stats([high_voice], speaker_id="high")
    assert s_low["sufficient"] and s_high["sufficient"]

    q_low = quantise_plan(low_voice, s_low)
    q_high = quantise_plan(high_voice, s_high)
    assert [w["pitch_bin"] for w in q_low] == [w["pitch_bin"] for w in q_high]
    assert q_low[-1]["pitch_bin"] == N_PITCH_BINS - 1

    pooled = fit_speaker_stats([low_voice, high_voice], speaker_id="pooled")
    assert pooled["sufficient"]
    p_low = quantise_plan(low_voice, pooled)
    p_high = quantise_plan(high_voice, pooled)
    assert p_low[-1]["pitch_bin"] < p_high[-1]["pitch_bin"], (
        p_low[-1]["pitch_bin"], p_high[-1]["pitch_bin"])


def test_sparse_speaker_falls_back_to_global_bins_and_is_flagged():
    n = MIN_WORDS_FOR_SPEAKER_BINS - 1
    words = _synthetic_words(n, rng_seed=3)
    stats = fit_speaker_stats([words], speaker_id="sparse")
    assert stats["sufficient"] is False
    assert str(n) in stats["reason"] and "min_words" in stats["reason"]

    q = quantise_plan(words, stats)
    assert {w["bin_source"] for w in q} == {"global"}


def test_degenerate_spread_is_reported_not_silently_binned():
    words = [{"word": f"w{i}", "dur_s": 0.3, "dur_frames": 6,
              "f0_mean_hz": 180.0, "energy_rms": 0.05} for i in range(200)]
    stats = fit_speaker_stats([words], speaker_id="flat")
    assert stats["sufficient"] is False
    assert "degenerate" in stats["reason"], stats["reason"]


def test_wrong_bin_count_in_speaker_stats_raises():
    words = _synthetic_words(120, rng_seed=4)
    stats = fit_speaker_stats([words], n_pitch_bins=8, speaker_id="odd")
    assert stats["sufficient"]
    with pytest.raises(PlanQuantiseError, match="token surface"):
        quantise_plan(words, stats)


def test_silent_word_gets_the_unvoiced_bin_and_no_nan_leaks():
    audio = np.concatenate([_harmonic(200.0, 0.8), _silence(0.4),
                            _harmonic(210.0, 0.8)])
    spans = _spans_from_durations([0.8, 0.4, 0.8], words=["hel", "sil", "lo"])
    feats = word_features_from_spans(audio, SR, spans)

    assert feats[1]["voiced_frac"] == 0.0
    assert math.isnan(feats[1]["f0_mean_hz"])
    assert feats[1]["silent"] is True
    assert not math.isnan(feats[0]["f0_mean_hz"])

    q = quantise_plan(feats)
    assert q[1]["pitch_bin"] == PITCH_UNVOICED_BIN
    assert q[1]["pitch_val"] is None
    assert q[1]["energy_bin"] == 0
    for w in q:
        assert not math.isnan(float(w["energy_val"]))
        assert isinstance(w["dur_bin"], int)

    toks = plan_to_tokens(q)
    assert "nan" not in toks.lower()
    assert toks.count(PLAN_WORD_TOKEN) == 3 and toks.count("[P:") == 2
    assert tokens_to_plan(toks)[1]["pitch_bin"] == PITCH_UNVOICED_BIN

    dq = dequantise_plan(q)
    assert dq[1]["pitch_hat"] is None


def test_unvoiced_words_do_not_enter_the_pitch_edges():
    words = _synthetic_words(120, rng_seed=5)
    voiced_only = [dict(w) for w in words[20:]]
    for w in words[:20]:
        w["f0_mean_hz"] = float("nan")
        w["voiced_frac"] = 0.0

    stats = fit_speaker_stats([words], speaker_id="mixed")
    assert stats["sufficient"] is True
    assert stats["n_words"] == 120 and stats["n_voiced_words"] == 100
    ref = fit_speaker_stats([voiced_only], speaker_id="voiced")
    assert stats["pitch_edges"] == pytest.approx(ref["pitch_edges"])

    q = quantise_plan(words, stats)
    assert all(w["pitch_bin"] == PITCH_UNVOICED_BIN for w in q[:20])


def test_plan_token_roundtrip_is_exact():
    words = _synthetic_words(60, rng_seed=6)
    words[3]["f0_mean_hz"] = float("nan")
    stats = fit_speaker_stats([words], speaker_id="spk")
    q = quantise_plan(words, stats)

    text = plan_to_tokens(q)
    back = tokens_to_plan(text)
    assert back == [{"pitch_bin": w["pitch_bin"], "dur_bin": w["dur_bin"],
                     "energy_bin": w["energy_bin"]} for w in q]
    assert plan_to_tokens(back) == text
    assert text.count(PLAN_WORD_TOKEN) == len(q)

    wrapped = plan_to_tokens(q, wrap=True)
    assert wrapped.startswith("<PLAN>") and wrapped.endswith("</PLAN>")
    assert tokens_to_plan(wrapped) == back


def test_plan_tokens_carry_no_separators():
    text = plan_to_tokens(quantise_plan(_synthetic_words(5, rng_seed=8)))
    assert " " not in text, text
    assert text.startswith(PLAN_WORD_TOKEN)
    spaced = text.replace("][", "] [")
    assert tokens_to_plan(spaced) == tokens_to_plan(text)


def test_plan_vocab_matches_the_configured_bin_counts():
    v = plan_vocab()
    assert len(v) == len(set(v)) == 3 + N_PITCH_BINS + 16 + N_ENERGY_BINS
    assert PLAN_WORD_TOKEN in v
    assert "[P:0]" in v and f"[P:{N_PITCH_BINS - 1}]" in v
    assert f"[P:{N_PITCH_BINS}]" not in v and "[P:U]" not in v


@pytest.mark.parametrize("bad", [
    "[P:3][D:2][E:1]",
    "[PW][P:3][E:1][D:2]",
    "[PW][P:3][D:2]",
    "[PW][P:3][D:2][E:99]",
    "[PW][P:99][D:2][E:1]",
    "[PW][P:3][D:U][E:1]",
    "[PW][P:3][D:2][E:1] garbage",
    "[PW][P][D:2][E:1]",
    "[PW][P:3][P:4][D:2][E:1]",
])


def test_malformed_plan_strings_raise(bad):
    with pytest.raises(PlanTokenError):
        tokens_to_plan(bad)


def test_explicit_unvoiced_symbol_parses_if_a_future_vocab_adds_it():
    assert tokens_to_plan("[PW][P:U][D:2][E:1]") == [
        {"pitch_bin": PITCH_UNVOICED_BIN, "dur_bin": 2, "energy_bin": 1}]


def test_config_plan_vocab_covers_every_token_we_emit():
    if verify_config_vocab() is None:
        pytest.skip("config.py has no plan tier yet")


def _voiced_utterance(dur_s=2.0):
    return _harmonic(180.0, dur_s)


def test_low_confidence_alignment_is_flagged_and_can_raise():
    audio = _voiced_utterance()

    def align(_a, words):
        step = 2.0 / len(words)
        return [(w, i * step, (i + 1) * step, 0.05) for i, w in enumerate(words)]

    ex = PlanExtractor(align_fn=align)
    res = ex.extract(audio, SR, "hello world")
    assert res.ok is False
    assert "low_conf" in res.flags, res.flags
    assert "mean_conf" in res.reason and "0.05" in res.reason
    assert res.mean_conf == pytest.approx(0.05)
    assert len(res.words) == 2

    with pytest.raises(PlanExtractError, match="low_conf"):
        ex.extract(audio, SR, "hello world", strict=True)


def test_one_weak_word_is_advisory_not_a_row_rejection():
    audio = _voiced_utterance()

    def align(_a, words):
        step = 2.0 / len(words)
        confs = [0.95] * len(words)
        confs[1] = 0.05
        return [(w, i * step, (i + 1) * step, confs[i])
                for i, w in enumerate(words)]

    res = PlanExtractor(align_fn=align).extract(audio, SR, "one two three four five")
    assert res.ok is True
    assert "low_conf_word" in res.flags
    assert res.words[1]["conf"] == pytest.approx(0.05)


def test_dropped_words_show_up_as_low_coverage():
    audio = _voiced_utterance()

    def align(_a, words):
        kept = words[:-1]
        step = 2.0 / max(len(kept), 1)
        return [(w, i * step, (i + 1) * step, 0.9) for i, w in enumerate(kept)]

    res = PlanExtractor(align_fn=align).extract(audio, SR, "one two three four")
    assert res.ok is False
    assert "low_coverage" in res.flags
    assert res.coverage == pytest.approx(0.75)


def test_alignment_exception_propagates_loudly():
    def align(_a, _w):
        raise PlanAlignmentError("emission too short")

    with pytest.raises(PlanAlignmentError, match="emission too short"):
        PlanExtractor(align_fn=align).extract(_voiced_utterance(), SR, "hi there")


def test_spans_past_the_end_of_the_audio_raise():
    def align(_a, words):
        return [(w, 2.0 * i, 2.0 * (i + 1), 0.9) for i, w in enumerate(words)]

    with pytest.raises(PlanAlignmentError, match="frame-to-second"):
        PlanExtractor(align_fn=align).extract(_voiced_utterance(2.0), SR, "one two")


def test_extract_word_plan_reuses_one_extractor_per_configuration():
    from plan_extract import _EXTRACTOR_CACHE

    def align(_a, words):
        return [(w, 0.5 * i, 0.5 * (i + 1), 0.9) for i, w in enumerate(words)]

    audio = _voiced_utterance()
    a = extract_word_plan(audio, SR, "one two", align_fn=align)
    b = extract_word_plan(audio, SR, "one two", align_fn=align)
    assert a.words == b.words
    assert sum(1 for k in _EXTRACTOR_CACHE if ("align_fn", align) in k) == 1

    with pytest.raises(TypeError, match="mutually exclusive"):
        extract_word_plan(audio, SR, "one two",
                          extractor=PlanExtractor(align_fn=align), align_fn=align)


def test_invalid_span_raises_rather_than_emitting_garbage_times():
    with pytest.raises(PlanExtractError, match="non-positive duration"):
        word_features_from_spans(_voiced_utterance(), SR, [("w", 0.5, 0.5)])
    with pytest.raises(PlanExtractError, match="non-positive duration"):
        word_features_from_spans(_voiced_utterance(), SR, [("w", 0.9, 0.4)])


def test_digital_silence_and_empty_transcript_raise():
    ex = PlanExtractor(align_fn=lambda a, w: [(w[0], 0.0, 1.0, 0.9)])
    with pytest.raises(PlanExtractError, match="silence"):
        ex.extract(_silence(2.0), SR, "hello")
    with pytest.raises(PlanExtractError, match="empty transcript"):
        ex.extract(_voiced_utterance(), SR, "   ")
    with pytest.raises(PlanExtractError, match="normalisation"):
        ex.extract(_voiced_utterance(), SR, "1234 5678")
    with pytest.raises(PlanExtractError, match="too short"):
        ex.extract(_harmonic(180.0, 0.02), SR, "hello")


def test_frame_rate_mismatch_between_stages_raises():
    feats = word_features_from_spans(_voiced_utterance(), SR, [("w", 0.0, 0.8)],
                                     frame_rate=12.5)
    assert feats[0]["dur_frames"] == 10
    with pytest.raises(PlanQuantiseError, match="frame rate"):
        quantise_plan(feats, frame_rate=FRAME_RATE_HZ)
    assert quantise_plan(feats, frame_rate=12.5)[0]["dur_frames"] == 10


def test_transcript_normalisation_keeps_apostrophes_and_reports_drops():
    words, dropped = normalise_transcript("Don't  stop, 42 — it's Fine!")
    assert words == ["don't", "stop", "it's", "fine"]
    assert dropped == ["42", "—"]


def test_extract_end_to_end_with_injected_aligner():
    audio = np.concatenate([_harmonic(180.0, 0.6), _silence(0.2),
                            _harmonic(300.0, 0.6)])

    def align(_a, words):
        return [(words[0], 0.0, 0.6, 0.95), (words[1], 0.8, 1.4, 0.9)]

    res = PlanExtractor(align_fn=align).extract(audio, SR, "Hello, world!")
    assert res.ok is True and res.flags == []
    assert res.coverage == 1.0
    assert res.duration_s == pytest.approx(1.4, abs=1e-3)
    assert res.time_coverage == pytest.approx(1.2 / 1.4, abs=1e-3)
    assert [w["word"] for w in res.words] == ["hello", "world"]
    assert abs(res.words[0]["f0_mean_hz"] - 180.0) / 180.0 < 0.03
    assert abs(res.words[1]["f0_mean_hz"] - 300.0) / 300.0 < 0.03
    assert list(res) == res.words and len(res) == 2


def test_extraction_and_quantisation_are_deterministic():
    audio = np.concatenate([_harmonic(190.0, 0.7), _harmonic(260.0, 0.7)])
    spans = _spans_from_durations([0.7, 0.7])

    a = word_features_from_spans(audio, SR, spans)
    b = word_features_from_spans(audio, SR, spans)
    assert json.dumps(a, sort_keys=True) == json.dumps(b, sort_keys=True)

    words = _synthetic_words(120, rng_seed=11)
    s1 = fit_speaker_stats([words], speaker_id="d")
    s2 = fit_speaker_stats([words], speaker_id="d")
    assert json.dumps(s1, sort_keys=True) == json.dumps(s2, sort_keys=True)
    assert plan_to_tokens(quantise_plan(words, s1)) == \
           plan_to_tokens(quantise_plan(words, s2))


def test_quantise_does_not_mutate_its_input():
    words = _synthetic_words(60, rng_seed=12)
    before = json.dumps(words, sort_keys=True)
    quantise_plan(words, fit_speaker_stats([words], speaker_id="m"))
    assert json.dumps(words, sort_keys=True) == before


def test_a_speaker_with_no_positive_energy_reports_insufficiency():
    words = _synthetic_words(120, rng_seed=5)
    for w in words:
        w["energy_rms"] = 0.0
    stats = fit_speaker_stats([words], speaker_id="silent")
    assert stats["sufficient"] is False
    assert "positive energy" in stats["reason"]
    assert stats["pitch_edges"] is None and stats["energy_edges"] is None

    assert fit_speaker_stats([_synthetic_words(120, rng_seed=6)],
                             speaker_id="ok")["sufficient"] is True


def test_every_duration_bin_is_reachable_by_an_integer():
    import numpy as np
    edges = np.log(np.asarray(DUR_EDGES_FRAMES, dtype=np.float64))
    reached = {int(np.searchsorted(edges, np.log(f))) for f in range(1, 201)}
    from plan_extract import _bin_index
    reached = {int(_bin_index(np.log(f), edges)) for f in range(1, 201)}
    n_bins = len(DUR_EDGES_FRAMES) + 1
    dead = sorted(set(range(n_bins)) - reached)
    assert not dead, (f"duration bins {dead} unreachable by any integer frame "
                      f"count — see Amendment 3; do not ship a grid with dead bins")