led-controller/tools/generate_beat_test_track.py

#!/usr/bin/env python3
"""
Metronome-style mono click track for testing the audio beat detector.

Without ``-o``: streams S16LE PCM to ``aplay`` (stdin) until you press Ctrl+C.

With ``-o``: writes a WAV file of fixed length and exits.

Examples:
  python3 tools/generate_beat_test_track.py
  python3 tools/generate_beat_test_track.py --bpm 90
  python3 tools/generate_beat_test_track.py -o tests/audio/beat_test_120bpm.wav --duration 30
"""

from __future__ import annotations

import argparse
import math
import shutil
import struct
import subprocess
import wave
from pathlib import Path


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument(
        "-o",
        "--output",
        type=Path,
        default=None,
        help="If set, write this WAV file and exit (no live playback)",
    )
    p.add_argument("--bpm", type=float, default=120.0, help="Beats per minute (default: 120)")
    p.add_argument(
        "--duration",
        type=float,
        default=30.0,
        help="With -o only: click section length in seconds after intro (default: 30)",
    )
    p.add_argument(
        "--intro-silence",
        type=float,
        default=0.5,
        help="Leading silence in seconds (default: 0.5)",
    )
    p.add_argument("--sample-rate", type=int, default=44100, help="Sample rate Hz (default: 44100)")
    p.add_argument(
        "--click-ms",
        type=float,
        default=18.0,
        help="Approximate click length in ms (default: 18)",
    )
    p.add_argument(
        "--freq",
        type=float,
        default=1000.0,
        help="Click sine frequency Hz (default: 1000)",
    )
    return p.parse_args()


def _click_int16_samples(sr: int, click_ms: float, freq: float) -> tuple[list[int], int]:
    """One click; returns samples and click_len (same as len(samples))."""
    click_len = max(1, int(sr * max(4.0, click_ms) / 1000.0))
    freq_clamped = max(200.0, min(4000.0, float(freq)))
    floats: list[float] = []
    for i in range(click_len):
        t = i / sr
        env = math.sin(0.5 * math.pi * (i + 1) / click_len) ** 2
        floats.append(env * math.sin(2.0 * math.pi * freq_clamped * t))
    peak = max(abs(x) for x in floats) or 1.0
    scale = 0.92 / peak
    out: list[int] = []
    for x in floats:
        v = int(round(max(-1.0, min(1.0, x * scale)) * 32767.0))
        out.append(max(-32767, min(32767, v)))
    return out, click_len


def _render_scaled_samples(
    sr: int,
    bpm: float,
    intro: float,
    dur: float,
    click_ms: float,
    freq: float,
) -> tuple[list[float], int, float, int]:
    beat_sec = 60.0 / bpm
    click_len = max(1, int(sr * max(4.0, click_ms) / 1000.0))
    freq_clamped = max(200.0, min(4000.0, float(freq)))
    total_sec = intro + dur
    n_samples = int(sr * total_sec)
    intro_samples = int(sr * intro)

    samples = [0.0] * n_samples
    beat_samples = int(round(sr * beat_sec))
    if beat_samples < click_len + 1:
        raise SystemExit("BPM too high for this sample rate / click length")

    beat_idx = 0
    while True:
        start = intro_samples + beat_idx * beat_samples
        if start >= n_samples:
            break
        for i in range(click_len):
            pos = start + i
            if pos >= n_samples:
                break
            t = i / sr
            env = math.sin(0.5 * math.pi * (i + 1) / click_len) ** 2
            s = env * math.sin(2.0 * math.pi * freq_clamped * t)
            samples[pos] += s
        beat_idx += 1

    peak = max(abs(x) for x in samples) or 1.0
    scale = 0.92 / peak
    for i in range(n_samples):
        samples[i] = max(-1.0, min(1.0, samples[i] * scale))

    return samples, sr, total_sec, beat_idx


def write_wav_mono16(path: Path, samples: list[float], sr: int) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with wave.open(str(path), "w") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sr)
        for x in samples:
            v = int(round(x * 32767.0))
            w.writeframes(struct.pack("<h", v))


def _stream_aplay_until_interrupt(
    sr: int,
    bpm: float,
    intro_silence: float,
    click_ms: float,
    freq: float,
) -> None:
    aplay = shutil.which("aplay")
    if not aplay:
        raise SystemExit("aplay not found; install alsa-utils, or use -o to write a WAV file.")

    click_samps, click_len = _click_int16_samples(sr, click_ms, freq)
    beat_samples = int(round(sr * 60.0 / bpm))
    if beat_samples < click_len + 1:
        raise SystemExit("BPM too high for this sample rate / click length")

    silence_samples = beat_samples - click_len
    beat_chunk = struct.pack("<" + "h" * len(click_samps), *click_samps) + (
        b"\x00\x00" * silence_samples
    )

    intro_samples = int(sr * max(0.0, float(intro_silence)))
    intro_chunk = b"\x00\x00" * intro_samples

    argv = [aplay, "-q", "-t", "raw", "-f", "S16_LE", "-c", "1", "-r", str(sr)]
    proc = subprocess.Popen(argv, stdin=subprocess.PIPE)
    if proc.stdin is None:
        raise SystemExit("aplay did not open stdin")

    print(
        f"Streaming {bpm} BPM, {sr} Hz mono -> aplay (raw). Ctrl+C to stop.",
        flush=True,
    )
    try:
        if intro_chunk:
            proc.stdin.write(intro_chunk)
        beats_written = 0
        # Write several beats per syscall to reduce overhead
        batch = 8
        multi = beat_chunk * batch
        while True:
            proc.stdin.write(multi)
            beats_written += batch
            if beats_written % 256 == 0:
                proc.stdin.flush()
    except BrokenPipeError:
        print("aplay exited.", flush=True)
    except KeyboardInterrupt:
        print("\nStopped.", flush=True)
    finally:
        try:
            proc.stdin.close()
        except BrokenPipeError:
            pass
        proc.wait(timeout=3)


def main() -> None:
    args = _parse_args()
    sr = max(8000, min(96000, int(args.sample_rate)))
    bpm = max(40.0, min(240.0, float(args.bpm)))

    if args.output is not None:
        intro = max(0.0, float(args.intro_silence))
        dur = max(1.0, float(args.duration))
        samples, sr_u, total_sec, beats = _render_scaled_samples(
            sr, bpm, intro, dur, float(args.click_ms), float(args.freq)
        )
        write_wav_mono16(args.output, samples, sr_u)
        print(
            f"Wrote {args.output} ({len(samples)} samples, {total_sec:.1f}s, {sr_u} Hz mono): "
            f"{bpm} BPM, ~{beats} beats"
        )
        return

    _stream_aplay_until_interrupt(
        sr,
        bpm,
        float(args.intro_silence),
        float(args.click_ms),
        float(args.freq),
    )


if __name__ == "__main__":
    main()