led-controller/tests/play_varying_click_track.py

#!/usr/bin/env python3
"""Play a click track with tempo variation for BPM detector testing.

Examples:
  python tests/play_varying_click_track.py --start-bpm 90 --end-bpm 150 --seconds 60
  python tests/play_varying_click_track.py --pattern steps --step-bpms 100,120,140,160 --step-seconds 8
"""

from __future__ import annotations

import argparse
import math
import time

import numpy as np
import sounddevice as sd


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Play varying-BPM click track")
    parser.add_argument("--device", default=None, help="Output audio device name or index")
    parser.add_argument("--sample-rate", type=int, default=44100, help="Output sample rate")
    parser.add_argument("--seconds", type=float, default=60.0, help="Playback duration")
    parser.add_argument(
        "--stabilize-seconds",
        type=float,
        default=5.0,
        help="Play initial BPM for this many seconds before varying",
    )
    parser.add_argument("--pattern", choices=("sweep", "steps"), default="sweep")

    # Sweep options
    parser.add_argument("--start-bpm", type=float, default=90.0, help="Sweep start BPM")
    parser.add_argument("--end-bpm", type=float, default=150.0, help="Sweep end BPM")

    # Step options
    parser.add_argument(
        "--step-bpms",
        default="100,120,140,160",
        help="Comma-separated BPMs for step mode",
    )
    parser.add_argument("--step-seconds", type=float, default=8.0, help="Seconds per BPM step")

    # Click tone options
    parser.add_argument("--click-ms", type=float, default=25.0, help="Click duration")
    parser.add_argument("--accent-every", type=int, default=4, help="Accent every N beats")
    parser.add_argument(
        "--hold-beats",
        type=int,
        default=1,
        help="Hold BPM for this many beats before recalculating",
    )
    return parser.parse_args()


def bpm_for_time(t_s: float, args: argparse.Namespace, step_bpms: list[float]) -> float:
    if t_s < args.stabilize_seconds:
        return args.start_bpm if args.pattern == "sweep" else (step_bpms[0] if step_bpms else 120.0)

    adj_t = t_s - args.stabilize_seconds
    if args.pattern == "sweep":
        active_seconds = max(1e-6, args.seconds - args.stabilize_seconds)
        if active_seconds <= 0:
            return args.start_bpm
        alpha = min(1.0, max(0.0, adj_t / active_seconds))
        return args.start_bpm + (args.end_bpm - args.start_bpm) * alpha

    if not step_bpms:
        return 120.0
    if args.step_seconds <= 0:
        return step_bpms[0]
    idx = int(adj_t // args.step_seconds) % len(step_bpms)
    return step_bpms[idx]


def main() -> int:
    args = parse_args()
    if args.seconds <= 0:
        raise SystemExit("--seconds must be > 0")
    if args.sample_rate <= 0:
        raise SystemExit("--sample-rate must be > 0")

    step_bpms = [float(x.strip()) for x in args.step_bpms.split(",") if x.strip()]
    click_samples = max(1, int(args.click_ms * args.sample_rate / 1000.0))

    state = {
        "t": 0.0,  # playback time in seconds
        "next_beat": 0.0,
        "beat_idx": 0,
        "current_bpm": 0.0,
        "held_bpm": 0.0,
        "hold_counter": 0,
        "click_remaining": 0,
        "click_phase": 0,
        "click_freq": 1320.0,
        "click_amp": 0.6,
    }

    def callback(outdata, frames, _time_info, status):
        if status:
            print(f"audio status: {status}")
        block = np.zeros(frames, dtype=np.float32)
        for i in range(frames):
            t = state["t"]
            dynamic_bpm = bpm_for_time(t, args, step_bpms)
            state["current_bpm"] = dynamic_bpm
            bpm = state["held_bpm"] if state["held_bpm"] > 0 else dynamic_bpm
            beat_interval = 60.0 / max(1e-6, bpm)

            if t >= state["next_beat"]:
                hold_beats = max(1, int(args.hold_beats))
                if state["hold_counter"] <= 0:
                    state["held_bpm"] = dynamic_bpm
                    state["hold_counter"] = hold_beats
                state["hold_counter"] -= 1
                bpm = state["held_bpm"]
                beat_interval = 60.0 / max(1e-6, bpm)
                state["beat_idx"] += 1
                state["next_beat"] = t + beat_interval
                state["click_remaining"] = click_samples
                state["click_phase"] = 0
                accented = (
                    args.accent_every > 0 and state["beat_idx"] % args.accent_every == 1
                )
                state["click_freq"] = 1760.0 if accented else 1320.0
                state["click_amp"] = 0.9 if accented else 0.6

            if state["click_remaining"] > 0:
                p = state["click_phase"]
                env = math.exp(-8.0 * (p / click_samples))
                sample = state["click_amp"] * env * math.sin(
                    2.0 * math.pi * state["click_freq"] * (p / args.sample_rate)
                )
                block[i] = sample
                state["click_phase"] += 1
                state["click_remaining"] -= 1

            state["t"] += 1.0 / args.sample_rate

        outdata[:, 0] = block

    print(
        f"Playing varying click track for {args.seconds:.1f}s ({args.pattern}), "
        f"stabilize={args.stabilize_seconds:.1f}s"
    )
    with sd.OutputStream(
        samplerate=args.sample_rate,
        channels=1,
        dtype="float32",
        callback=callback,
        device=args.device,
        blocksize=0,
    ):
        start = time.time()
        last_printed_beat = 0
        while (time.time() - start) < args.seconds:
            beat_idx = int(state["beat_idx"])
            if beat_idx != last_printed_beat:
                last_printed_beat = beat_idx
                print(
                    f"beat={beat_idx:04d} bpm={state['held_bpm']:.2f} "
                    f"(target={state['current_bpm']:.2f})"
                )
            time.sleep(0.05)

    print("Done.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())