#!/usr/bin/env python3
"""
Record amplifier output via sound-card line-in and estimate THD+N of a 1 kHz tone.

WARNING — SAFETY / EQUIPMENT PROTECTION
---------------------------------------
Speaker-level amplifier outputs can be TENS OF VOLTS. A PC sound-card line input
expects roughly 1 Vrms max. You MUST use a resistive attenuator / voltage divider
(and preferably a protection clamp) between the amp output and the sound card.

Example divider (VERIFY power rating and loading): for ~40 Vpk max into line-in
aim ~1 Vpk → attenuation ≈ 40:1, e.g. 39 kΩ series + 1 kΩ to GND across the input.
Use non-inductive resistors rated for the fault power. Never connect speaker
outputs directly to a sound card.

Usage:
  # Analyze an existing recording of a 1 kHz tone through the attenuator:
  python measure_thd.py --wav recording_1k.wav --f0 1000

  # Or record live (requires sounddevice + working input device):
  python measure_thd.py --record 5 --device 0 --f0 1000
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
from scipy.io import wavfile
from scipy.signal import windows


def load_wav_mono(path: Path) -> tuple[int, np.ndarray]:
    fs, data = wavfile.read(str(path))
    if data.ndim > 1:
        data = data.mean(axis=1)
    if np.issubdtype(data.dtype, np.integer):
        max_int = np.iinfo(data.dtype).max
        data = data.astype(np.float64) / max_int
    else:
        data = data.astype(np.float64)
    return int(fs), data


def record_audio(seconds: float, fs: int, device: int | None) -> np.ndarray:
    try:
        import sounddevice as sd
    except ImportError as exc:
        raise SystemExit(
            "sounddevice is required for --record. pip install sounddevice\n"
            "Or pass --wav path/to/recording.wav instead."
        ) from exc
    print(
        "Recording… ensure attenuator is connected. "
        "Speaker-level → sound card without a divider can destroy the input."
    )
    audio = sd.rec(
        int(seconds * fs),
        samplerate=fs,
        channels=1,
        dtype="float64",
        device=device,
    )
    sd.wait()
    return audio[:, 0]


def thd_n(
    x: np.ndarray,
    fs: int,
    f0: float,
    n_harmonics: int = 10,
) -> dict[str, float]:
    """
    FFT-based THD and time/FFT hybrid THD+N.

    THD  = sqrt(sum harmonic_powers) / fundamental_amplitude
    THD+N = rms(signal with fundamental notched) / rms(fundamental)

    A narrow spectral notch removes the fundamental (±bins) from a copy of
    the spectrum; the residual RMS over the audio band is distortion+noise.
    """
    n = len(x)
    if n < fs:
        raise ValueError("need at least 1 second of audio")
    x = x[n // 10 : -n // 10] if n > 2 * fs else x
    x = x - np.mean(x)

    win = windows.hann(len(x), sym=False)
    xw = x * win
    cg = np.sum(win) / len(win)
    spec = np.fft.rfft(xw)
    freqs = np.fft.rfftfreq(len(xw), d=1.0 / fs)
    # Peak amplitude spectrum (coherent-gain corrected)
    mag = np.abs(spec) / (len(xw) * cg)

    def peak_amp(freq: float, half_width_hz: float = 3.0) -> float:
        mask = (freqs >= freq - half_width_hz) & (freqs <= freq + half_width_hz)
        if not np.any(mask):
            return 0.0
        return float(np.max(mag[mask]))

    fund_a = peak_amp(f0)
    if fund_a <= 0:
        raise RuntimeError(f"no fundamental found near {f0} Hz — check recording")

    harm_ssq = 0.0
    for k in range(2, n_harmonics + 1):
        fk = k * f0
        if fk >= fs / 2:
            break
        ha = peak_amp(fk)
        harm_ssq += ha * ha

    thd = np.sqrt(harm_ssq) / fund_a

    # THD+N via time-domain least-squares fundamental removal
    t = np.arange(len(x)) / fs
    c = np.cos(2 * np.pi * f0 * t)
    s = np.sin(2 * np.pi * f0 * t)
    a_c = 2 * np.dot(x, c) / len(x)
    a_s = 2 * np.dot(x, s) / len(x)
    fund_rms = np.sqrt(0.5 * (a_c**2 + a_s**2))
    resid_td = x - (a_c * c + a_s * s)
    resid_rms = np.sqrt(np.mean(resid_td**2))
    thd_n_ratio = resid_rms / (fund_rms + 1e-20)

    return {
        "thd_ratio": float(thd),
        "thd_percent": float(100.0 * thd),
        "thd_n_ratio": float(thd_n_ratio),
        "thd_n_percent": float(100.0 * thd_n_ratio),
        "thd_n_db": float(20.0 * np.log10(thd_n_ratio + 1e-20)),
        "fundamental_amp": fund_a,
    }


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("--wav", type=Path, help="Existing mono/stereo WAV of 1 kHz tone")
    p.add_argument("--record", type=float, metavar="SEC", help="Record SEC seconds")
    p.add_argument("--fs", type=int, default=48000)
    p.add_argument("--device", type=int, default=None, help="sounddevice input index")
    p.add_argument("--f0", type=float, default=1000.0, help="Fundamental frequency")
    p.add_argument("--save-recording", type=Path, default=None)
    return p.parse_args()


def main() -> int:
    args = parse_args()
    if args.wav:
        fs, x = load_wav_mono(args.wav)
    elif args.record:
        fs = args.fs
        x = record_audio(args.record, fs, args.device)
        if args.save_recording:
            from scipy.io.wavfile import write

            write(
                str(args.save_recording),
                fs,
                (np.clip(x, -1, 1) * 32767).astype(np.int16),
            )
            print(f"saved {args.save_recording}")
    else:
        print("Provide --wav FILE or --record SECONDS", file=sys.stderr)
        return 2

    results = thd_n(x, fs, args.f0)
    print(f"Sample rate: {fs} Hz, samples: {len(x)}")
    print(f"THD:    {results['thd_percent']:.4f} %")
    print(f"THD+N:  {results['thd_n_percent']:.4f} %  ({results['thd_n_db']:.2f} dB)")
    print(
        "Note: absolute accuracy depends on sound-card noise, attenuator loading, "
        "and that the amp is not clipping."
    )
    return 0


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