src/audio.py

"""Audio helpers: WAV I/O, resampling, augmentation and synthetic noise.

Kept dependency-light: only numpy is required (no ffmpeg / scipy), so it
runs anywhere including a minimal Hugging Face Space.
"""

from __future__ import annotations

import hashlib
import io
import math
import random
import re
import wave
from pathlib import Path
from typing import Tuple

import numpy as np


# --------------------------------------------------------------------------- #
# Naming helpers
# --------------------------------------------------------------------------- #

def slugify(text: str, max_len: int = 80) -> str:
    text = str(text).strip().lower()
    text = re.sub(r"[^a-z0-9]+", "_", text)
    text = re.sub(r"_+", "_", text).strip("_")
    return (text or "item")[:max_len]


def stable_hash(text: str, length: int = 12) -> str:
    return hashlib.sha1(text.encode("utf-8")).hexdigest()[:length]


# --------------------------------------------------------------------------- #
# WAV read / write / resample
# --------------------------------------------------------------------------- #

def read_wav_bytes(wav_bytes: bytes) -> Tuple[np.ndarray, int]:
    """Decode 16-bit PCM WAV bytes into mono float32 samples and sample rate."""
    with wave.open(io.BytesIO(wav_bytes), "rb") as wf:
        channels = wf.getnchannels()
        sample_width = wf.getsampwidth()
        sample_rate = wf.getframerate()
        frames = wf.readframes(wf.getnframes())

    if sample_width != 2:
        raise ValueError(f"Expected 16-bit PCM WAV, got sample width {sample_width}")

    audio = np.frombuffer(frames, dtype=np.int16).astype(np.float32)
    if channels > 1:
        audio = audio.reshape(-1, channels).mean(axis=1)
    return audio, sample_rate


def read_wav_file(path: Path) -> Tuple[np.ndarray, int]:
    return read_wav_bytes(Path(path).read_bytes())


def write_wav_file(path: Path, audio: np.ndarray, sample_rate_hz: int) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    audio_i16 = np.clip(audio, -32768, 32767).astype(np.int16)
    with wave.open(str(path), "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate_hz)
        wf.writeframes(audio_i16.tobytes())


def resample(audio: np.ndarray, src_rate: int, dst_rate: int) -> np.ndarray:
    """Linear-interpolation resample. Good enough for a bootstrap dataset."""
    if src_rate == dst_rate or len(audio) == 0:
        return audio.astype(np.float32)
    duration = len(audio) / float(src_rate)
    dst_len = max(1, int(round(duration * dst_rate)))
    src_idx = np.linspace(0.0, len(audio) - 1, num=dst_len)
    return np.interp(src_idx, np.arange(len(audio)), audio).astype(np.float32)


# --------------------------------------------------------------------------- #
# Shaping
# --------------------------------------------------------------------------- #

def normalize(audio: np.ndarray, peak: float = 28000.0) -> np.ndarray:
    if len(audio) == 0:
        return audio
    m = float(np.max(np.abs(audio)))
    if m < 1.0:
        return audio
    return (audio * (peak / m)).astype(np.float32)


def pad_or_trim(audio: np.ndarray, target_samples: int, random_crop: bool = False) -> np.ndarray:
    current = len(audio)
    if current == target_samples:
        return audio
    if current > target_samples:
        start = (
            random.randint(0, current - target_samples)
            if random_crop
            else (current - target_samples) // 2
        )
        return audio[start:start + target_samples]
    pad_total = target_samples - current
    pad_left = pad_total // 2
    pad_right = pad_total - pad_left
    return np.pad(audio, (pad_left, pad_right), mode="constant")


# --------------------------------------------------------------------------- #
# Augmentation
# --------------------------------------------------------------------------- #

def _gain(audio: np.ndarray, gain_db: float) -> np.ndarray:
    return audio * (10.0 ** (gain_db / 20.0))


def _time_shift(audio: np.ndarray, max_shift: int) -> np.ndarray:
    return np.roll(audio, random.randint(-max_shift, max_shift))


def _add_noise(audio: np.ndarray, snr_db: float) -> np.ndarray:
    noise = np.random.normal(0.0, 1.0, len(audio)).astype(np.float32)
    clean_power = float(np.mean(audio ** 2))
    noise_power = float(np.mean(noise ** 2))
    if clean_power < 1.0 or noise_power < 1e-9:
        return audio
    target_noise_power = clean_power / (10.0 ** (snr_db / 10.0))
    noise *= math.sqrt(target_noise_power / noise_power)
    return audio + noise


def _echo(audio: np.ndarray, sr: int) -> np.ndarray:
    delay = random.randint(int(0.03 * sr), int(0.12 * sr))
    decay = random.uniform(0.08, 0.25)
    out = audio.copy()
    if 0 < delay < len(audio):
        out[delay:] += audio[:-delay] * decay
    return out


def augment(audio: np.ndarray, sr: int) -> np.ndarray:
    out = audio.copy()
    out = _gain(out, random.uniform(-6.0, 3.0))
    out = _time_shift(out, max_shift=int(0.12 * sr))
    if random.random() < 0.75:
        out = _add_noise(out, random.choice([30, 25, 20, 15, 10]))
    if random.random() < 0.35:
        out = _echo(out, sr)
    return normalize(out, 28000.0)


# --------------------------------------------------------------------------- #
# Synthetic background noise
# --------------------------------------------------------------------------- #

def _white(n: int) -> np.ndarray:
    return np.random.normal(0.0, 1.0, n).astype(np.float32)


def _pink(n: int) -> np.ndarray:
    white = _white(n)
    out = np.zeros_like(white)
    alpha = 0.985
    for i in range(1, n):
        out[i] = alpha * out[i - 1] + (1.0 - alpha) * white[i]
    return out.astype(np.float32)


def _brown(n: int) -> np.ndarray:
    brown = np.cumsum(_white(n))
    brown = brown - np.mean(brown)
    return normalize(brown.astype(np.float32), 1.0)


def _hum(n: int, sr: int) -> np.ndarray:
    t = np.arange(n, dtype=np.float32) / float(sr)
    hum = (
        np.sin(2.0 * math.pi * 50.0 * t)
        + 0.5 * np.sin(2.0 * math.pi * 100.0 * t)
        + 0.25 * np.sin(2.0 * math.pi * 150.0 * t)
    )
    hum += 0.04 * _white(n)
    return hum.astype(np.float32)


def _fan(n: int, sr: int) -> np.ndarray:
    base = _pink(n)
    t = np.arange(n, dtype=np.float32) / float(sr)
    blade_rate = random.uniform(18.0, 45.0)
    modulation = 0.65 + 0.35 * np.sin(2.0 * math.pi * blade_rate * t)
    return (base * modulation).astype(np.float32)


def _cafe(n: int, sr: int) -> np.ndarray:
    base = 0.55 * _pink(n) + 0.45 * _white(n)
    transient_count = max(1, int((n / sr) * random.uniform(2.0, 6.0)))
    for _ in range(transient_count):
        pos = random.randint(0, max(0, n - 1))
        length = random.randint(max(1, int(0.008 * sr)), max(2, int(0.05 * sr)))
        end = min(n, pos + length)
        if end <= pos:
            continue
        click = np.hanning(end - pos).astype(np.float32)
        base[pos:end] += click * random.uniform(0.5, 2.0)
    return base.astype(np.float32)


def _street(n: int, sr: int) -> np.ndarray:
    base = 0.7 * _brown(n) + 0.3 * _white(n)
    t = np.arange(n, dtype=np.float32) / float(sr)
    for _ in range(random.randint(1, 3)):
        center = random.uniform(0.2, max(0.21, t[-1] - 0.2))
        width = random.uniform(0.2, 0.8)
        envelope = np.exp(-0.5 * ((t - center) / width) ** 2)
        freq = random.uniform(70.0, 180.0)
        base += 0.35 * envelope * np.sin(2.0 * math.pi * freq * t)
    return base.astype(np.float32)


def make_background_noise(noise_type: str, num_samples: int, sr: int) -> np.ndarray:
    if noise_type == "white":
        noise = _white(num_samples)
    elif noise_type == "pink":
        noise = _pink(num_samples)
    elif noise_type == "brown":
        noise = _brown(num_samples)
    elif noise_type == "hum":
        noise = _hum(num_samples, sr)
    elif noise_type == "fan":
        noise = _fan(num_samples, sr)
    elif noise_type == "cafe":
        noise = _cafe(num_samples, sr)
    elif noise_type == "street":
        noise = _street(num_samples, sr)
    else:
        raise ValueError(f"Unknown noise type: {noise_type}")
    return normalize(noise, random.uniform(6000, 22000))