scene-actor-extraction/scripts/validation/test_audio_offset.py

#!/usr/bin/env python3
"""
VR-014 — the v1 audio signature recovers a known trim offset on real audio.

TRACES: UT-105, UT-106, UT-107, UT-108 | VR-014 | IR-004

    python scripts/validation/test_audio_offset.py [build_dir]

The golden vector (IR-005) proves the *arithmetic* is identical in both
producers. It cannot prove the thing the signature exists for: that when the
same cut arrives trimmed differently, sliding one signature against the other
finds the true alignment and only the true alignment. Its fixture is a synthetic
tone sweep, which is pathologically easy to align; film dialogue and score are
not, and that is what this measures.

The signature is computed by the **shipped C++**, through the `sae_audio`
nanobind module — never a numpy port. A third implementation of a fingerprint
whose whole value rests on three implementations agreeing byte for byte would be
the one nobody checks against the golden vector.

The slide *is* written here in numpy, deliberately: matching is the consumer's
algorithm (server SPEC.md section 3), owned by the server and the jRay plugin,
not by this repo. Writing it out is what makes this a test of the signature
rather than a test of somebody's matcher.

Two independent offset mechanisms are checked, because they can fail
separately:

  * a **window offset** (UT-105) — two 120 s excerpts taken from different
    points, which is the alignment search itself; and
  * a **head trim** (UT-106) — a real file with delta seconds removed from the
    front, which additionally exercises the runtime/2 anchor: the window follows
    the midpoint, so cutting delta from the head moves it by delta/2, not delta.
    That factor of two is the easiest thing in the whole feature to get wrong
    and nothing else checks it.
"""

import base64
import random
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

import numpy as np

REPO = Path(__file__).resolve().parent.parent.parent
BUILD = Path(sys.argv[1]) if len(sys.argv) > 1 else REPO / "build"
sys.path.insert(0, str(BUILD))

import sae_audio  # noqa: E402

FIXTURE = REPO / "tests" / "fixtures" / "audio" / "superhero_offset_200s.flac"
TONE = REPO / "tests" / "fixtures" / "audio" / "jray_audio_v1_tone.flac"

# Server SPEC.md section 3, "Matching and offset recovery". The cap is the
# spec's, not a convenience: +/-600 frames is ~56 s, which covers realistic trim
# differences, and an offset outside it must be declined rather than guessed at.
SEARCH_CAP_FRAMES = 600
AUDIO_TIER = 0.85
LOOSE_TIER = 0.60

TRIALS = 40
SEED = 20250731

HOP_SEC = sae_audio.hop_size / sae_audio.sample_rate

# What the offset is actually *for*: shifting scene windows, which are seconds
# long. Half a second of error is invisible against them, and that budget is
# what makes the numbers below readable — an offset is quantised to whole
# frames, so no correct answer can be worse than half a frame (46 ms) and the
# feature has an order of magnitude in hand before anything is at stake.
OFFSET_BUDGET_SEC = 0.5


def peak_bins(signature):
    """The per-frame peak band index, which is what the slide compares.

    The `v1:` prefix is checked against the constant the C++ exports rather
    than a literal, so a producer bump cannot be silently parsed as v1 here
    (IR-008).
    """
    prefix = sae_audio.version_prefix
    if not signature.startswith(prefix):
        raise AssertionError(f"signature is not {prefix!r}: {signature[:8]!r}")
    packed = np.frombuffer(base64.b64decode(signature[len(prefix):]), dtype=np.uint8)
    if np.any(packed & 0x80):
        raise AssertionError("reserved bit set — not a structurally valid signature")
    return packed >> 2


def best_match(reference, query, cap=SEARCH_CAP_FRAMES, slack=0):
    """Slide `query` against `reference`; return (score, offset_frames).

    `offset` is how many frames later the query's window begins, so
    ``query[i]`` lines up with ``reference[i + offset]``. Score is the fraction
    of overlapping frames whose peak bin agrees, exactly as the spec defines it.

    `slack` widens what counts as agreement to a frame within +/-slack, which is
    not the spec's rule — it is the candidate remedy UT-108 measures. It changes
    the *score* only; the offset it reports is still a whole-frame alignment.
    """
    best_score, best_offset = -1.0, 0
    for offset in range(-cap, cap + 1):
        if offset >= 0:
            a, b = reference[offset:], query[: len(query) - offset]
        else:
            a, b = reference[: len(reference) + offset], query[-offset:]
        n = min(len(a), len(b))
        if n < 100:            # too little overlap to mean anything
            continue
        a, b = a[:n], b[:n]
        if slack == 0:
            agree = a == b
        else:
            agree = np.zeros(n, dtype=bool)
            for shift in range(-slack, slack + 1):
                shifted = np.roll(a, shift)
                # 255 is not a band index, so the wrapped end can never agree.
                if shift > 0:
                    shifted[:shift] = 255
                elif shift < 0:
                    shifted[shift:] = 255
                agree |= shifted == b
        score = float(np.mean(agree))
        if score > best_score:
            best_score, best_offset = score, offset
    return best_score, best_offset


def decode_mono(path):
    """The whole fixture as float32 mono at 11025 Hz — the signature's own rate."""
    raw = subprocess.run(
        ["ffmpeg", "-nostdin", "-v", "error", "-i", str(path),
         "-ac", "1", "-ar", str(sae_audio.sample_rate), "-f", "f32le", "-"],
        capture_output=True, check=True).stdout
    return np.frombuffer(raw, dtype="<f4")


def trim_head(source, seconds, out):
    """`source` with `seconds` removed from the front — a differently trimmed release."""
    subprocess.run(
        ["ffmpeg", "-nostdin", "-v", "error", "-y", "-ss", f"{seconds:.3f}", "-i", str(source),
         "-ac", "1", "-ar", str(sae_audio.sample_rate), "-sample_fmt", "s16",
         "-c:a", "flac", str(out)], check=True)
    return out


def tier(score):
    if score >= AUDIO_TIER:
        return "audio"
    return "loose" if score >= LOOSE_TIER else "none"


# ── The random trial set, signed once and reused ─────────────────────────────

def random_trials(pcm):
    """(reference bins, [(expected_frames, query bins)]) for TRIALS excerpts.

    Signing 40 windows is the expensive part of this file, so UT-105 and UT-108
    share one set — they ask different questions of the same measurements.
    """
    window = sae_audio.window_samples
    reference = peak_bins(sae_audio.signature_from_mono(pcm[:window]))
    rng = random.Random(SEED)
    queries = []

    for _ in range(TRIALS):
        # Within the search cap: past it, no offset is recoverable by
        # construction, which UT-107 checks separately.
        start = rng.randrange(0, SEARCH_CAP_FRAMES * sae_audio.hop_size)
        signature = sae_audio.signature_from_mono(pcm[start:start + window])
        assert signature is not None, "a full window must always sign"
        queries.append((start / sae_audio.hop_size, peak_bins(signature)))

    return reference, queries


# ── UT-105 — window offsets from random excerpt starts ───────────────────────

def test_random_window_offsets(reference, queries):
    """Every in-cap offset is recovered to the nearest frame, on real audio."""
    rows = []
    for want, query in queries:
        score, offset = best_match(reference, query)
        rows.append((want, offset, score, abs(want - round(want))))

    expected = np.array([r[0] for r in rows])
    offset = np.array([r[1] for r in rows])
    score = np.array([r[2] for r in rows])
    subframe = np.array([r[3] for r in rows])
    error = np.abs(offset - expected)

    # The offset is quantised to whole frames, so the best any correct answer
    # can do is half a frame — 46 ms. What matters is the budget that half-frame
    # is measured against, and it is an order of magnitude away from it.
    assert error.max() <= 1.0, f"offset missed by {error.max():.2f} frames"
    assert error.max() * HOP_SEC <= OFFSET_BUDGET_SEC, (
        f"offset error {error.max() * HOP_SEC:.3f}s exceeds the {OFFSET_BUDGET_SEC}s budget")
    # Never mistaken for different content. This is the floor that matters: the
    # audio genuinely is the same cut, so a "no match" would be a false negative
    # on the case the feature exists for.
    assert score.min() >= LOOSE_TIER, f"same content scored {score.min():.3f}"
    # An offset that lands near a frame boundary has no excuse: it should reach
    # the top tier, and does.
    aligned = subframe <= 0.1
    assert aligned.any(), "seed no longer produces a near-aligned trial"
    assert score[aligned].min() >= AUDIO_TIER, (
        f"near-frame-aligned offset scored only {score[aligned].min():.3f}")

    print(f"UT-105  {TRIALS} random window offsets, all within the +/-600 frame cap")
    print(f"        offset error   : max {error.max():.2f} frames"
          f" = {error.max() * HOP_SEC * 1000:.0f} ms, against a"
          f" {OFFSET_BUDGET_SEC * 1000:.0f} ms budget")
    print(f"        score          : min {score.min():.3f}  median {np.median(score):.3f}"
          f"  max {score.max():.3f}")
    print("        score by sub-frame misalignment — the offset is exact in every row:")
    for lo, hi in ((0.0, 0.1), (0.1, 0.2), (0.2, 0.3), (0.3, 0.4), (0.4, 0.5)):
        m = (subframe >= lo) & (subframe < hi)
        if m.any():
            print(f"          {lo:.1f}-{hi:.1f} frame  n={m.sum():2d}"
                  f"  score {score[m].min():.3f}-{score[m].max():.3f}"
                  f"  tier {tier(np.median(score[m]))}")
    counts = {t: int(sum(1 for s in score if tier(s) == t)) for t in ("audio", "loose", "none")}
    print(f"        tiers          : {counts}")
    return counts


# ── UT-106 — head trims through real files, including the runtime/2 anchor ───

def test_head_trims():
    """A release with delta seconds of head removed aligns at delta/2 frames."""
    reference = peak_bins(sae_audio.compute_signature(str(FIXTURE)))
    results = []

    with tempfile.TemporaryDirectory() as tmp:
        for delta in (7.0, 23.5, 41.25, 60.0):
            trimmed = trim_head(FIXTURE, delta, Path(tmp) / f"trim_{delta}.flac")
            signature = sae_audio.compute_signature(str(trimmed))
            assert signature is not None, f"trim of {delta}s should still sign"
            score, offset = best_match(reference, peak_bins(signature))
            # The window follows the midpoint, so removing delta from the head
            # moves it by delta/2 — not by delta.
            expected = (delta / 2.0) / HOP_SEC
            assert abs(offset - expected) <= 1.0, (
                f"head trim {delta}s: expected ~{expected:.1f} frames, got {offset}")
            assert score >= LOOSE_TIER, f"head trim {delta}s scored {score:.3f}"
            results.append((delta, expected, offset, score))

    print("UT-106  head trims through the real decode path (compute_signature on a file)")
    for delta, expected, offset, score in results:
        print(f"        -{delta:6.2f}s head  expected {expected:7.2f} fr"
              f"  recovered {offset:5d}  score {score:.3f} ({tier(score)})")


# ── UT-107 — what must NOT match ─────────────────────────────────────────────

def test_declines(pcm, reference):
    """Out-of-cap offsets and unrelated content are declined, not guessed at."""
    window = sae_audio.window_samples

    beyond = int(75.0 * sae_audio.sample_rate)          # ~807 frames, past the cap
    assert beyond + window <= len(pcm), "fixture too short for the out-of-cap case"
    far = peak_bins(sae_audio.signature_from_mono(pcm[beyond:beyond + window]))
    score_beyond, offset_beyond = best_match(reference, far)
    assert score_beyond < LOOSE_TIER, (
        f"an offset past the cap scored {score_beyond:.3f} at {offset_beyond} — "
        "the search invented an alignment rather than declining")

    tone = peak_bins(sae_audio.compute_signature(str(TONE)))
    score_tone, offset_tone = best_match(reference, tone)
    assert score_tone < LOOSE_TIER, f"unrelated content scored {score_tone:.3f}"

    print("UT-107  declines rather than guesses")
    print(f"        offset past the +/-600 frame cap : best {score_beyond:.3f}"
          f" at {offset_beyond} ({tier(score_beyond)})")
    print(f"        unrelated content (tone fixture) : best {score_tone:.3f}"
          f" at {offset_tone} ({tier(score_tone)})")
    return far, tone


# ── UT-108 — the sub-frame demotion, and what one frame of slack costs ───────

def test_scoring_slack(reference, queries, far, tone):
    """Measured: +/-1 frame of slack in the *score* restores the `audio` tier.

    UT-105 leaves a real question open. Every offset is right, but two thirds of
    them score below the server's 0.85 `audio` threshold purely because the two
    windows' frame grids do not coincide — so a correctly aligned release is
    demoted to `loose`, which is the tier meaning "possibly the same cut,
    degraded audio". The obvious remedy is to stop demanding that frames line up
    exactly, and the question is what that costs in discrimination.

    Nothing is asserted about the spec's own rule here; this measures a
    candidate change to it, which is the server's to make (SPEC.md section 3).
    """
    print("UT-108  cost of relaxing the score's frame alignment")
    print(f"        {'slack':>5} {'audio':>6} {'loose':>6} {'none':>5}"
          f" {'min true':>9} {'worst err':>10} {'unrelated':>10} {'out-of-cap':>11}")

    measured = {}
    for slack in (0, 1, 2):
        score, error = [], []
        for want, query in queries:
            s, offset = best_match(reference, query, slack=slack)
            score.append(s)
            error.append(abs(offset - want))
        score, error = np.array(score), np.array(error)
        false_tone, _ = best_match(reference, tone, slack=slack)
        false_far, _ = best_match(reference, far, slack=slack)
        counts = {t: int(sum(1 for s in score if tier(s) == t)) for t in ("audio", "loose", "none")}
        measured[slack] = (score, error, max(false_tone, false_far))
        print(f"        {slack:>5} {counts['audio']:>6} {counts['loose']:>6} {counts['none']:>5}"
              f" {score.min():>9.3f} {error.max() * HOP_SEC * 1000:>7.0f} ms"
              f" {false_tone:>10.3f} {false_far:>11.3f}")

    score, error, worst_false = measured[1]
    # One frame of slack lifts every correct alignment to the top tier...
    assert score.min() >= AUDIO_TIER, (
        f"one frame of slack still leaves a true match at {score.min():.3f}")
    # ...without narrowing the gap that makes the threshold mean anything...
    assert worst_false < LOOSE_TIER, (
        f"slack lifted a false match to {worst_false:.3f}")
    # ...and the offset it costs is still far inside the budget: the score's
    # peak flattens slightly, so the argmax can pick an adjacent frame.
    assert error.max() * HOP_SEC <= OFFSET_BUDGET_SEC, (
        f"slack cost {error.max() * HOP_SEC:.3f}s of offset accuracy")
    print(f"        +/-1 frame: every true match reaches `audio` (min {score.min():.3f}),"
          f" worst false stays at {worst_false:.3f},")
    print(f"        and the offset costs {error.max() * HOP_SEC * 1000:.0f} ms of a"
          f" {OFFSET_BUDGET_SEC * 1000:.0f} ms budget. +/-2 buys nothing more.")


def main():
    if not FIXTURE.exists():
        print(f"missing fixture {FIXTURE} — regenerate with make_offset_fixture.sh", file=sys.stderr)
        return 2
    if shutil.which("ffmpeg") is None:
        print("this validation needs the ffmpeg CLI to trim the fixture", file=sys.stderr)
        return 2

    print(f"VR-014  audio-signature offset recovery on {FIXTURE.name}")
    print(f"        {sae_audio.expected_frames} frames per signature,"
          f" {HOP_SEC * 1000:.2f} ms per frame, cap +/-{SEARCH_CAP_FRAMES} frames")

    pcm = decode_mono(FIXTURE)
    reference, queries = random_trials(pcm)
    counts = test_random_window_offsets(reference, queries)
    test_head_trims()
    far, tone = test_declines(pcm, reference)
    test_scoring_slack(reference, queries, far, tone)

    print()
    print(f"PASS — every in-cap offset recovered to the nearest frame, worst"
          f" {1000 * HOP_SEC / 2:.0f} ms against a {OFFSET_BUDGET_SEC * 1000:.0f} ms budget.")
    if counts["audio"] < TRIALS:
        # Stated rather than asserted against the spec's rule: the offset is
        # right in every case, so this is the 0.85 threshold meeting a sub-frame
        # shift, not a defect in the signature. The threshold was calibrated on
        # a re-encode at zero offset, where the score is 1.00. UT-108 measures
        # the remedy; adopting it is the server spec's call, not this repo's.
        print(f"NOTE  — under the spec's exact-frame score only {counts['audio']}/{TRIALS}"
              f" reach `audio`; {counts['loose']} are demoted to `loose` by sub-frame"
              " shift alone. See UT-108.")
    return 0


if __name__ == "__main__":
    sys.exit(main())