ml/data/augment_min.py

"""
ml/data/augment_min.py - Minimal, real-only audio augmentation
===============================================================
Augmentation is applied ONLY to REAL recordings (gold corpora / user
clips). It never invents a disorder. Purpose: make the model robust to
real-world recording conditions (noise, device bandwidth, loudness,
tempo) — the point is that a speaker is not penalised for *recording
quality*, only for speech content.

Bounded magnitude + small probability so label meaning is never flipped
(a lisp stays a lisp; a stutter stays a stutter).

Transforms:
    add_noise    : low SNR floor (room/mic hiss)
    random_gain  : global loudness scaling
    time_stretch : slight tempo change, pitch-preserving (WSOLA)
    highpass     : remove low rumble / mic thump
"""
from __future__ import annotations
import numpy as np
import soundfile as sf
from pathlib import Path
from typing import Optional

import librosa
from scipy import signal as sp_signal


def add_noise(audio: np.ndarray, snr_db: float = 22.0, rng: Optional[np.random.Generator] = None) -> np.ndarray:
    rng = rng or np.random.default_rng()
    power = np.mean(audio ** 2) + 1e-12
    npow = power / (10 ** (snr_db / 10.0))
    noise = rng.normal(0.0, np.sqrt(npow), audio.shape).astype(np.float32)
    return audio + noise


def random_gain(audio: np.ndarray, factor: float = 1.0) -> np.ndarray:
    return audio * float(factor)


def random_speed(audio: np.ndarray, sr: int, factor: float = 1.0) -> np.ndarray:
    """Tempo change preserving pitch; bound factor ~0.97-1.04."""
    return librosa.effects.time_stretch(audio.astype(np.float64), rate=1.0 / factor).astype(np.float32)


def highpass(audio: np.ndarray, sr: int, cutoff: float = 70.0) -> np.ndarray:
    sos = sp_signal.butter(1, cutoff, "hp", fs=sr, output="sos")
    return sp_signal.sosfilt(sos, audio).astype(np.float32)


def augment(
    audio: np.ndarray,
    sr: int,
    *,
    prob_gain: float = 0.3,
    prob_noise: float = 0.4,
    prob_speed: float = 0.3,
    prob_hp: float = 0.2,
    rng: Optional[np.random.Generator] = None,
) -> np.ndarray:
    """Bounded random subset of real-world transforms. Preserves content."""
    rng = rng or np.random.default_rng()
    x = audio.astype(np.float32).copy()
    if rng.random() < prob_gain:
        x = random_gain(x, rng.uniform(0.7, 1.3))
    if rng.random() < prob_noise:
        x = add_noise(x, rng.uniform(18, 28), rng)
    if rng.random() < prob_speed:
        x = random_speed(x, sr, rng.uniform(0.97, 1.04))
    if rng.random() < prob_hp:
        x = highpass(x, sr)
    return np.clip(x, -0.99, 0.99).astype(np.float32)


def augment_wav(src: Path, dst: Path, sr: int = 16000, seed: Optional[int] = None) -> None:
    """Augment a wav file; deterministic if seed given."""
    y, _ = sf.read(str(src), dtype="float32")
    rng = np.random.default_rng(seed)
    out = augment(y, sr, rng=rng)
    sf.write(str(dst), out, sr, subtype="PCM_16")


if __name__ == "__main__":
    # self-check: run each transform, confirm length/validity preserved
    sr = 16000
    t = np.linspace(0, 1, sr, endpoint=False)
    tone = (0.5 * np.sin(2 * np.pi * 220 * t)).astype(np.float32)
    checks = {
        "noise":        lambda: add_noise(tone, 20, np.random.default_rng(1)),
        "gain":         lambda: random_gain(tone, 1.2),
        "speed_stretch":lambda: random_speed(tone, sr, 1.03),
        "highpass":     lambda: highpass(tone, sr),
        "augment_full": lambda: augment(tone, sr),
    }
    ok = True
    for name, fn in checks.items():
        try:
            out = fn()
            assert len(out) > 0 and np.isfinite(out.all()), f"{name} produced non-finite/empty"
            print(f"{name:15} OK len={len(out)}")
        except Exception as e:
            ok = False
            print(f"{name:15} FAIL {e}")
    raise SystemExit(0 if ok else 1)