drumsynth.py

"""Tiny stdlib drum synthesizer for Drum School lesson audio.

Renders each drum piece as a one-shot sample (synthesized once, cached),
then mixes the one-shots into a full-length track at the hit times from a
drum_tab-style hit list. No numpy, no external deps — `math`, `random`,
`array`, `wave` only, so it runs on the server's stock Python.

The goal is a clear, friendly play-along guide track (kids hear what they
should be playing), not a sample-accurate acoustic kit.
"""

from __future__ import annotations

import math
import random
import wave
from array import array
from pathlib import Path

SR = 44100


# ── DSP helpers ───────────────────────────────────────────────────────────────

def _noise(n: int, seed: int) -> list[float]:
    rng = random.Random(seed)
    return [rng.uniform(-1.0, 1.0) for _ in range(n)]


def _highpass(samples: list[float], alpha: float) -> list[float]:
    """One-pole high-pass. alpha in (0,1); higher keeps more highs."""
    out = [0.0] * len(samples)
    prev_in = 0.0
    prev_out = 0.0
    for i, x in enumerate(samples):
        prev_out = alpha * (prev_out + x - prev_in)
        prev_in = x
        out[i] = prev_out
    return out


def _lowpass(samples: list[float], alpha: float) -> list[float]:
    """One-pole low-pass. alpha in (0,1); higher keeps more highs."""
    out = [0.0] * len(samples)
    prev = 0.0
    for i, x in enumerate(samples):
        prev += alpha * (x - prev)
        out[i] = prev
    return out


def _sweep(f_start: float, f_end: float, dur: float, decay: float,
           drive: float = 0.0) -> list[float]:
    """Pitch-swept sine with exponential decay. `drive` adds soft clipping."""
    n = int(SR * dur)
    out = [0.0] * n
    phase = 0.0
    for i in range(n):
        t = i / SR
        f = f_start + (f_end - f_start) * min(1.0, t / max(dur * 0.6, 1e-6))
        phase += 2.0 * math.pi * f / SR
        s = math.sin(phase) * math.exp(-t * decay)
        if drive > 0.0:
            s = math.tanh(s * (1.0 + drive))
        out[i] = s
    return out


def _tone(freqs: list[float], dur: float, decay: float) -> list[float]:
    n = int(SR * dur)
    out = [0.0] * n
    scale = 1.0 / max(1, len(freqs))
    for i in range(n):
        t = i / SR
        env = math.exp(-t * decay)
        out[i] = sum(math.sin(2.0 * math.pi * f * t) for f in freqs) * scale * env
    return out


def _noise_hit(dur: float, decay: float, seed: int,
               hp: float = 0.0, lp: float = 0.0) -> list[float]:
    n = int(SR * dur)
    s = _noise(n, seed)
    if hp > 0.0:
        s = _highpass(s, hp)
    if lp > 0.0:
        s = _lowpass(s, lp)
    for i in range(n):
        s[i] *= math.exp(-(i / SR) * decay)
    return s


def _mix_oneshot(*parts: tuple[list[float], float]) -> list[float]:
    n = max(len(p) for p, _ in parts)
    out = [0.0] * n
    for p, gain in parts:
        for i, v in enumerate(p):
            out[i] += v * gain
    peak = max(1e-9, max(abs(v) for v in out))
    if peak > 1.0:
        out = [v / peak for v in out]
    return out


# ── One-shot bank ─────────────────────────────────────────────────────────────
# Deterministic seeds so regenerated packs are byte-identical.

def _build_oneshot(piece: str) -> list[float]:
    if piece == "kick":
        return _mix_oneshot(
            (_sweep(115.0, 44.0, 0.30, 16.0, drive=1.2), 1.0),
            (_noise_hit(0.012, 90.0, seed=11, hp=0.5), 0.35),
        )
    if piece == "snare":
        return _mix_oneshot(
            (_tone([182.0, 268.0], 0.22, 26.0), 0.55),
            (_noise_hit(0.26, 21.0, seed=21, hp=0.35), 0.85),
        )
    if piece == "snare_xstick":
        return _mix_oneshot(
            (_tone([760.0, 1210.0], 0.07, 60.0), 0.6),
            (_noise_hit(0.05, 80.0, seed=22, hp=0.6), 0.4),
        )
    if piece == "hh_closed":
        return _noise_hit(0.09, 55.0, seed=31, hp=0.82)
    if piece == "hh_pedal":
        return _mix_oneshot((_noise_hit(0.07, 70.0, seed=32, hp=0.75), 0.55))
    if piece == "hh_open":
        return _noise_hit(0.55, 6.5, seed=33, hp=0.8)
    if piece == "tom_hi":
        return _sweep(196.0, 128.0, 0.35, 11.0, drive=0.6)
    if piece == "tom_mid":
        return _sweep(152.0, 99.0, 0.40, 10.0, drive=0.6)
    if piece == "tom_low":
        return _sweep(120.0, 78.0, 0.45, 9.0, drive=0.6)
    if piece == "tom_floor":
        return _sweep(96.0, 58.0, 0.55, 8.0, drive=0.7)
    if piece in ("crash_l", "crash_r"):
        seed = 41 if piece == "crash_l" else 42
        return _mix_oneshot(
            (_noise_hit(1.6, 2.4, seed=seed, hp=0.7), 0.9),
            (_tone([523.0, 1244.0], 0.5, 6.0), 0.12),
        )
    if piece == "splash":
        return _noise_hit(0.55, 6.0, seed=43, hp=0.78)
    if piece == "china":
        return _mix_oneshot(
            (_noise_hit(1.0, 3.2, seed=44, hp=0.6), 1.0),
            (_noise_hit(1.0, 3.2, seed=45, hp=0.3, lp=0.5), 0.5),
        )
    if piece == "stack":
        return _noise_hit(0.16, 22.0, seed=46, hp=0.7)
    if piece == "ride":
        return _mix_oneshot(
            (_noise_hit(0.9, 4.2, seed=51, hp=0.75), 0.4),
            (_tone([515.0, 782.0], 0.7, 5.0), 0.28),
        )
    if piece == "ride_bell":
        return _tone([618.0, 934.0, 1246.0], 0.45, 8.0)
    if piece == "bell":
        return _tone([872.0, 1310.0], 0.42, 9.0)
    # Unknown piece — short neutral tick so a future vocabulary addition
    # still produces audible feedback rather than silence.
    return _noise_hit(0.06, 60.0, seed=99, hp=0.6)


_ONESHOT_CACHE: dict[str, list[float]] = {}


def oneshot(piece: str) -> list[float]:
    cached = _ONESHOT_CACHE.get(piece)
    if cached is None:
        cached = _build_oneshot(piece)
        _ONESHOT_CACHE[piece] = cached
    return cached


# Relative loudness per piece in the mix (post velocity scaling).
_PIECE_GAIN = {
    "kick": 1.0,
    "snare": 0.9,
    "snare_xstick": 0.65,
    "hh_closed": 0.42,
    "hh_pedal": 0.32,
    "hh_open": 0.5,
    "tom_hi": 0.85,
    "tom_mid": 0.85,
    "tom_low": 0.9,
    "tom_floor": 0.9,
    "crash_l": 0.75,
    "crash_r": 0.75,
    "splash": 0.6,
    "china": 0.7,
    "stack": 0.6,
    "ride": 0.5,
    "ride_bell": 0.6,
    "bell": 0.55,
}


# ── Track rendering ───────────────────────────────────────────────────────────

def render_track(hits: list[dict], duration_s: float) -> array:
    """Mix drum_tab-style hits into a mono float track, return int16 samples.

    Each hit: {"t": seconds, "p": piece-id, "v": 1-127 (optional),
    "k": choke-tail seconds (optional)}.
    """
    n = int(SR * (duration_s + 2.0))  # +2s tail so the last cymbal rings out
    buf = [0.0] * n
    for hit in hits:
        piece = hit.get("p", "")
        shot = oneshot(piece)
        vel = float(hit.get("v", 100)) / 127.0
        gain = _PIECE_GAIN.get(piece, 0.7) * (0.25 + 0.75 * vel)
        start = int(float(hit["t"]) * SR)
        if start >= n:
            continue
        choke = hit.get("k")
        length = len(shot)
        if isinstance(choke, (int, float)) and choke > 0:
            length = min(length, int((float(choke) + 0.03) * SR))
        end = min(n, start + length)
        fade_from = length - int(0.03 * SR) if length < len(shot) else length
        for i in range(end - start):
            v = shot[i]
            if i >= fade_from:  # quick fade for choked cymbals
                v *= max(0.0, 1.0 - (i - fade_from) / max(1, int(0.03 * SR)))
            buf[start + i] += v * gain
    peak = max(1e-9, max(abs(v) for v in buf))
    norm = min(1.0, 0.89 / peak)
    pcm = array("h", (int(max(-1.0, min(1.0, v * norm)) * 32767) for v in buf))
    return pcm


def write_wav(path: Path, pcm: array) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with wave.open(str(path), "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(SR)
        w.writeframes(pcm.tobytes())