#!/usr/bin/env python3
"""
Generate WAV test signals for amplifier / active-speaker validation.

Signals:
  - 1 kHz sine (THD stimulus)
  - Pink noise (ear / spectrum checks)
  - Log sine sweep 20 Hz–20 kHz (frequency-response deconvolution)

Examples:
  python generate_test_signals.py
  python generate_test_signals.py --outdir ./wav --fs 48000 --duration 10
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
from scipy.io import wavfile


def write_wav(path: Path, fs: int, data: np.ndarray) -> None:
    """Write float32 mono −1..1 as 16-bit PCM WAV."""
    clipped = np.clip(data, -1.0, 1.0)
    pcm = (clipped * 32767.0).astype(np.int16)
    wavfile.write(str(path), fs, pcm)
    print(f"wrote {path}  ({len(data) / fs:.2f}s @ {fs} Hz)")


def sine_1k(fs: int, duration: float, amplitude: float = 0.5) -> np.ndarray:
    """1 kHz sine — primary THD+N stimulus (leave headroom vs 0 dBFS)."""
    t = np.arange(int(fs * duration)) / fs
    return amplitude * np.sin(2 * np.pi * 1000.0 * t)


def pink_noise(fs: int, duration: float, amplitude: float = 0.3) -> np.ndarray:
    """
    Approximate pink noise via frequency-domain 1/f shaping of white noise.
    WHY pink: equal energy per octave — useful for listening / RTA checks.
    """
    n = int(fs * duration)
    rng = np.random.default_rng(42)
    white = rng.standard_normal(n)
    spec = np.fft.rfft(white)
    freqs = np.fft.rfftfreq(n, d=1.0 / fs)
    # Avoid div-by-zero at DC; shape ~ 1/sqrt(f) for pink amplitude spectrum
    scale = np.ones_like(freqs)
    nz = freqs > 0
    scale[nz] = 1.0 / np.sqrt(freqs[nz])
    pink = np.fft.irfft(spec * scale, n=n)
    pink /= np.max(np.abs(pink)) + 1e-12
    return amplitude * pink


def log_sine_sweep(
    fs: int,
    duration: float,
    f0: float = 20.0,
    f1: float = 20000.0,
    amplitude: float = 0.5,
) -> np.ndarray:
    """
    Exponential (log) sine sweep — Farina method friendly.
    WHY log sweep: long energy at LF, invertible for IR / frequency response.
    """
    n = int(fs * duration)
    t = np.arange(n) / fs
    # Instantaneous phase for exponential sweep
    if f0 <= 0 or f1 <= f0:
        raise ValueError("need 0 < f0 < f1")
    ln_ratio = np.log(f1 / f0)
    phase = 2 * np.pi * f0 * duration / ln_ratio * (np.exp(t * ln_ratio / duration) - 1.0)
    sweep = amplitude * np.sin(phase)
    # Short fade in/out to reduce clicks
    fade = int(0.01 * fs)
    if fade > 1:
        w = np.linspace(0, 1, fade)
        sweep[:fade] *= w
        sweep[-fade:] *= w[::-1]
    return sweep


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--outdir", type=Path, default=Path("wav_out"))
    p.add_argument("--fs", type=int, default=48000, help="Sample rate (Hz)")
    p.add_argument("--duration", type=float, default=5.0, help="Length seconds")
    p.add_argument("--amplitude", type=float, default=0.5, help="Peak amplitude < 1")
    return p.parse_args()


def main() -> int:
    args = parse_args()
    args.outdir.mkdir(parents=True, exist_ok=True)
    fs = args.fs
    dur = args.duration
    a = args.amplitude

    write_wav(args.outdir / "sine_1kHz.wav", fs, sine_1k(fs, dur, a))
    write_wav(args.outdir / "pink_noise.wav", fs, pink_noise(fs, dur, min(a, 0.3)))
    write_wav(
        args.outdir / "sine_sweep_20_20k.wav",
        fs,
        log_sine_sweep(fs, max(dur, 10.0), 20.0, 20000.0, a),
    )
    print("Done. Play through the amp chain at safe levels; never clip the DUT.")
    return 0


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