#!/usr/bin/env python3
"""
Plot frequency response from a recorded log sine sweep (and optional inverse sweep).

Preferred workflow (Farina / exponential sweep):
  1. Play measurement/wav_out/sine_sweep_20_20k.wav through the DUT.
  2. Record the attenuated amp output to recorded_sweep.wav.
  3. Run:
       python plot_frequency_response.py \\
           --recorded recorded_sweep.wav \\
           --reference ../measurement/wav_out/sine_sweep_20_20k.wav

If only a stepped-sine or noise recording is available, pass --recorded alone
with --mode psd for a rough spectrum (not a calibrated FR).

WARNING: Use a resistive attenuator before the sound-card input — see measure_thd.py.
"""

from __future__ import annotations

import argparse
from pathlib import Path

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


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


def frequency_response_from_sweep(
    recorded: np.ndarray,
    reference: np.ndarray,
    fs: int,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Estimate H(f) ≈ FFT(recorded) / FFT(reference) with regularization.
    Aligns signals via cross-correlation first.
    """
    n = min(len(recorded), len(reference))
    rec = recorded[:n] - np.mean(recorded[:n])
    ref = reference[:n] - np.mean(reference[:n])

    # Align
    corr = correlate(rec, ref, mode="full")
    lag = int(np.argmax(corr) - (len(ref) - 1))
    if lag > 0:
        rec = rec[lag:]
        ref = ref[: len(rec)]
    elif lag < 0:
        ref = ref[-lag:]
        rec = rec[: len(ref)]
    n = min(len(rec), len(ref))
    rec, ref = rec[:n], ref[:n]

    # Power-of-two FFT length
    nfft = 1 << int(np.ceil(np.log2(n)))
    win = windows.hann(n, sym=False)
    REC = np.fft.rfft(rec * win, n=nfft)
    REF = np.fft.rfft(ref * win, n=nfft)
    eps = 1e-12 * np.max(np.abs(REF))
    H = REC / (REF + eps)
    freqs = np.fft.rfftfreq(nfft, d=1.0 / fs)
    mag_db = 20.0 * np.log10(np.abs(H) + 1e-20)
    return freqs, mag_db


def rough_psd_db(x: np.ndarray, fs: int) -> tuple[np.ndarray, np.ndarray]:
    nfft = min(len(x), 65536)
    win = windows.hann(nfft, sym=False)
    seg = x[:nfft] * win
    spec = np.fft.rfft(seg)
    freqs = np.fft.rfftfreq(nfft, d=1.0 / fs)
    mag_db = 20.0 * np.log10(np.abs(spec) / np.max(np.abs(spec)) + 1e-20)
    return freqs, mag_db


def plot(freqs: np.ndarray, mag_db: np.ndarray, out: Path, title: str) -> None:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    mask = (freqs >= 20) & (freqs <= 20000)
    plt.figure(figsize=(10, 5))
    plt.semilogx(freqs[mask], mag_db[mask])
    plt.xlim(20, 20000)
    plt.xlabel("Frequency (Hz)")
    plt.ylabel("Magnitude (dB)")
    plt.title(title)
    plt.grid(True, which="both", ls=":")
    plt.tight_layout()
    plt.savefig(out, dpi=150)
    print(f"saved plot {out}")


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("--recorded", type=Path, required=True)
    p.add_argument("--reference", type=Path, default=None, help="Played sweep WAV")
    p.add_argument(
        "--mode",
        choices=("sweep", "psd"),
        default="sweep",
        help="sweep needs --reference; psd is relative spectrum only",
    )
    p.add_argument("--out", type=Path, default=Path("frequency_response.png"))
    return p.parse_args()


def main() -> int:
    args = parse_args()
    fs_r, rec = load_mono(args.recorded)

    if args.mode == "psd" or args.reference is None:
        freqs, mag = rough_psd_db(rec, fs_r)
        title = "Relative spectrum (PSD mode)"
    else:
        fs_ref, ref = load_mono(args.reference)
        if fs_ref != fs_r:
            raise SystemExit(f"sample rate mismatch: recorded {fs_r} vs reference {fs_ref}")
        freqs, mag = frequency_response_from_sweep(rec, ref, fs_r)
        # Normalize midband to 0 dB for readability
        mid = (freqs >= 500) & (freqs <= 2000)
        if np.any(mid):
            mag = mag - np.median(mag[mid])
        title = "Estimated frequency response (sweep deconvolution)"

    plot(freqs, mag, args.out, title)
    # Also dump CSV for KiCad/docs / further analysis
    csv_path = args.out.with_suffix(".csv")
    np.savetxt(
        csv_path,
        np.column_stack([freqs, mag]),
        delimiter=",",
        header="freq_hz,mag_db",
        comments="",
    )
    print(f"saved data {csv_path}")
    return 0


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