#!/usr/bin/env python3
"""Laptop-mic pitch analyzer for K1 accompaniment register checks.

Uses autocorrelation F0 estimate (numpy only + sounddevice).
Helps verify bass/chord octave when MIDI channel logs lack ch8(bass).

Examples:
  python tools/mic_pitch_analyze.py --list
  python tools/mic_pitch_analyze.py --seconds 8 --out mic_i.log
  python tools/mic_pitch_analyze.py --live
  python tools/mic_pitch_analyze.py --compare mic_i.wav mic_iii.wav
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
import wave
from collections import Counter
from dataclasses import dataclass
from datetime import datetime

import numpy as np

try:
    import sounddevice as sd
except ImportError as exc:
    raise SystemExit("Missing sounddevice: pip install sounddevice") from exc

A4_HZ = 440.0
A4_MIDI = 69


@dataclass
class PitchFrame:
    t: float
    hz: float
    midi: float
    note: str
    conf: float
    rms: float


NOTE_NAMES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]


def hz_to_midi(hz: float) -> float:
    if hz <= 0:
        return float("nan")
    return A4_MIDI + 12.0 * np.log2(hz / A4_HZ)


def midi_to_name(midi: float) -> str:
    if not np.isfinite(midi):
        return "--"
    n = int(round(midi))
    return f"{NOTE_NAMES[n % 12]}{n // 12 - 1}"


def list_devices() -> None:
    if hasattr(sys.stdout, "reconfigure"):
        try:
            sys.stdout.reconfigure(encoding="utf-8", errors="replace")
        except Exception:
            pass
    for i, d in enumerate(sd.query_devices()):
        name = str(d.get("name", "")).encode("utf-8", "replace").decode("utf-8", "replace")
        print(
            f"{i}: in={d['max_input_channels']} out={d['max_output_channels']} "
            f"sr={d['default_samplerate']} {name}",
            flush=True,
        )
    try:
        print("default device pair:", sd.default.device, flush=True)
    except Exception:
        pass


def record(seconds: float, sr: int, device: int | None) -> np.ndarray:
    frames = int(seconds * sr)
    print(f"Recording {seconds:.1f}s @ {sr} Hz (Ctrl+C to abort)...", flush=True)
    audio = sd.rec(frames, samplerate=sr, channels=1, dtype="float32", device=device)
    sd.wait()
    return audio[:, 0]


def save_wav(path: str, audio: np.ndarray, sr: int) -> None:
    pcm = np.clip(audio, -1.0, 1.0)
    pcm16 = (pcm * 32767.0).astype(np.int16)
    with wave.open(path, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sr)
        w.writeframes(pcm16.tobytes())


def load_wav(path: str) -> tuple[np.ndarray, int]:
    with wave.open(path, "rb") as w:
        sr = w.getframerate()
        nch = w.getnchannels()
        raw = w.readframes(w.getnframes())
        pcm = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
        if nch > 1:
            pcm = pcm.reshape(-1, nch).mean(axis=1)
    return pcm, sr


def estimate_f0_acorr(
    frame: np.ndarray,
    sr: int,
    fmin: float,
    fmax: float,
) -> tuple[float, float]:
    """Return (hz, confidence). confidence in [0,1] from normalized peak."""
    x = frame.astype(np.float64)
    x = x - np.mean(x)
    rms = float(np.sqrt(np.mean(x * x)) + 1e-12)
    if rms < 1e-4:
        return 0.0, 0.0

    # Hamming window
    x = x * np.hamming(len(x))
    corr = np.correlate(x, x, mode="full")
    corr = corr[len(corr) // 2 :]

    i_min = max(1, int(sr / fmax))
    i_max = min(len(corr) - 1, int(sr / fmin))
    if i_max <= i_min:
        return 0.0, 0.0

    seg = corr[i_min : i_max + 1]
    peak_rel = int(np.argmax(seg))
    peak = peak_rel + i_min
    if corr[0] <= 1e-12:
        return 0.0, 0.0
    conf = float(corr[peak] / corr[0])
    if conf < 0.25:
        return 0.0, conf

    # parabolic interpolation around peak
    if 1 <= peak < len(corr) - 1:
        a, b, c = corr[peak - 1], corr[peak], corr[peak + 1]
        denom = a - 2 * b + c
        if abs(denom) > 1e-12:
            peak = peak + 0.5 * (a - c) / denom

    hz = float(sr / peak)
    if hz < fmin or hz > fmax:
        return 0.0, conf
    return hz, conf


def analyze_audio(
    audio: np.ndarray,
    sr: int,
    hop_ms: float = 50.0,
    win_ms: float = 80.0,
    fmin: float = 40.0,
    fmax: float = 600.0,
    conf_min: float = 0.35,
    rms_min: float = 0.01,
) -> list[PitchFrame]:
    hop = max(1, int(sr * hop_ms / 1000.0))
    win = max(hop, int(sr * win_ms / 1000.0))
    out: list[PitchFrame] = []
    if len(audio) < win:
        return out

    for start in range(0, len(audio) - win, hop):
        frame = audio[start : start + win]
        rms = float(np.sqrt(np.mean(frame.astype(np.float64) ** 2)))
        hz, conf = estimate_f0_acorr(frame, sr, fmin, fmax)
        t = start / sr
        if hz <= 0 or conf < conf_min or rms < rms_min:
            out.append(PitchFrame(t, 0.0, float("nan"), "--", conf, rms))
            continue
        midi = hz_to_midi(hz)
        out.append(PitchFrame(t, hz, midi, midi_to_name(midi), conf, rms))
    return out


def summarize(frames: list[PitchFrame], label: str = "") -> dict:
    voiced = [f for f in frames if f.hz > 0 and np.isfinite(f.midi)]
    if not voiced:
        return {"label": label, "voiced": 0, "total": len(frames)}

    midis = np.array([f.midi for f in voiced], dtype=np.float64)
    hz = np.array([f.hz for f in voiced], dtype=np.float64)
    notes = [f.note for f in voiced]
    # round to nearest MIDI for histogram
    rounded = [int(round(m)) for m in midis]
    top = Counter(rounded).most_common(8)
    # bass-ish: MIDI <= 48 (C3)
    bass_ratio = float(np.mean(midis <= 48.0))
    low_ratio = float(np.mean(midis <= 40.0))  # <= E2
    return {
        "label": label,
        "voiced": len(voiced),
        "total": len(frames),
        "hz_median": float(np.median(hz)),
        "hz_p10": float(np.percentile(hz, 10)),
        "hz_p90": float(np.percentile(hz, 90)),
        "midi_median": float(np.median(midis)),
        "midi_p10": float(np.percentile(midis, 10)),
        "midi_p90": float(np.percentile(midis, 90)),
        "note_median": midi_to_name(float(np.median(midis))),
        "top_notes": [(midi_to_name(float(n)), c) for n, c in top],
        "bass_le_C3_ratio": bass_ratio,
        "low_le_E2_ratio": low_ratio,
    }


def print_summary(s: dict) -> None:
    if s.get("voiced", 0) == 0:
        print(f"[{s.get('label','')}] no pitched frames", flush=True)
        return
    print(
        f"[{s.get('label','')}] voiced={s['voiced']}/{s['total']}  "
        f"median={s['note_median']} ({s['midi_median']:.1f} / {s['hz_median']:.1f}Hz)  "
        f"p10={s['midi_p10']:.1f} p90={s['midi_p90']:.1f}  "
        f"<=E2={s['low_le_E2_ratio']*100:.0f}%  <=C3={s['bass_le_C3_ratio']*100:.0f}%",
        flush=True,
    )
    tops = ", ".join(f"{n}×{c}" for n, c in s["top_notes"][:5])
    print(f"  top: {tops}", flush=True)


def write_frame_log(path: str, frames: list[PitchFrame], summary: dict) -> None:
    with open(path, "w", encoding="utf-8", newline="\n") as f:
        f.write(f"# mic pitch {datetime.now().isoformat(timespec='seconds')}\n")
        f.write("# " + json.dumps(summary, ensure_ascii=False) + "\n")
        f.write("t_s,hz,midi,note,conf,rms\n")
        for fr in frames:
            midi = "" if not np.isfinite(fr.midi) else f"{fr.midi:.2f}"
            f.write(
                f"{fr.t:.3f},{fr.hz:.2f},{midi},{fr.note},{fr.conf:.3f},{fr.rms:.4f}\n"
            )


def live_monitor(
    sr: int,
    device: int | None,
    fmin: float,
    fmax: float,
    seconds: float,
) -> None:
    win = int(sr * 0.08)
    hop = int(sr * 0.05)
    print("Live pitch (Ctrl+C stop)...", flush=True)
    buf = np.zeros(0, dtype=np.float32)
    t0 = time.time()

    def callback(indata, frames, time_info, status):  # noqa: ARG001
        nonlocal buf
        if status:
            print(status, flush=True)
        buf = np.concatenate([buf, indata[:, 0].copy()])
        while len(buf) >= win:
            frame = buf[:win]
            buf = buf[hop:]
            hz, conf = estimate_f0_acorr(frame, sr, fmin, fmax)
            rms = float(np.sqrt(np.mean(frame.astype(np.float64) ** 2)))
            if hz > 0 and conf >= 0.35 and rms >= 0.01:
                midi = hz_to_midi(hz)
                print(
                    f"{time.time()-t0:6.1f}s  {midi_to_name(midi):4s}  "
                    f"midi={midi:5.1f}  {hz:6.1f}Hz  conf={conf:.2f}  rms={rms:.3f}",
                    flush=True,
                )

    with sd.InputStream(samplerate=sr, channels=1, dtype="float32", device=device, callback=callback):
        if seconds > 0:
            sd.sleep(int(seconds * 1000))
        else:
            while True:
                sd.sleep(200)


def compare_summaries(a: dict, b: dict) -> None:
    print("\n======== COMPARE ========", flush=True)
    print_summary(a)
    print_summary(b)
    if a.get("voiced", 0) == 0 or b.get("voiced", 0) == 0:
        print("Need pitched content in both takes.", flush=True)
        return
    d_midi = b["midi_median"] - a["midi_median"]
    print(f"median delta (B-A): {d_midi:+.2f} semitones", flush=True)
    if d_midi <= -9:
        print("PASS-ish: B is about an octave lower than A (expected III+ vs I/II).", flush=True)
    elif d_midi <= -5:
        print("PARTIAL: B lower than A but less than full octave.", flush=True)
    elif abs(d_midi) < 2:
        print("FAIL-ish: medians similar — register fix may not be audible on mic mix.", flush=True)
    else:
        print("CHECK: unexpected direction/amount; inspect top notes / low_le_E2 ratios.", flush=True)
    print(
        f"low<=E2: A={a['low_le_E2_ratio']*100:.0f}%  B={b['low_le_E2_ratio']*100:.0f}%",
        flush=True,
    )


def main() -> None:
    if hasattr(sys.stdout, "reconfigure"):
        try:
            sys.stdout.reconfigure(encoding="utf-8", errors="replace")
        except Exception:
            pass
    p = argparse.ArgumentParser(description="Mic pitch analyzer for K1 register tests")
    p.add_argument("--list", action="store_true", help="List audio devices")
    p.add_argument("--device", type=int, default=None, help="Input device index")
    p.add_argument("--sr", type=int, default=16000)
    p.add_argument("--seconds", type=float, default=8.0)
    p.add_argument("--live", action="store_true")
    p.add_argument("--wav", default="", help="Analyze existing wav instead of recording")
    p.add_argument("--out", default="", help="Prefix for wav/csv outputs")
    p.add_argument("--fmin", type=float, default=40.0, help="Min F0 Hz (bass ~41=E1)")
    p.add_argument("--fmax", type=float, default=600.0, help="Max F0 Hz")
    p.add_argument(
        "--compare",
        nargs=2,
        metavar=("A", "B"),
        help="Compare two wav/csv summary sources (wav preferred)",
    )
    p.add_argument("--label", default="")
    args = p.parse_args()

    if args.list:
        list_devices()
        return

    if args.compare:
        summaries = []
        for i, path in enumerate(args.compare):
            label = "A" if i == 0 else "B"
            if path.lower().endswith(".wav"):
                audio, sr = load_wav(path)
                frames = analyze_audio(audio, sr, fmin=args.fmin, fmax=args.fmax)
                s = summarize(frames, label=f"{label}:{os.path.basename(path)}")
            else:
                raise SystemExit("compare expects .wav files")
            summaries.append(s)
        compare_summaries(summaries[0], summaries[1])
        return

    if args.live:
        live_monitor(args.sr, args.device, args.fmin, args.fmax, args.seconds if args.seconds > 0 else 0)
        return

    prefix = args.out or f"mic_pitch_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    if args.wav:
        audio, sr = load_wav(args.wav)
        label = args.label or os.path.basename(args.wav)
    else:
        audio = record(args.seconds, args.sr, args.device)
        sr = args.sr
        label = args.label or "rec"
        wav_path = prefix if prefix.lower().endswith(".wav") else prefix + ".wav"
        save_wav(wav_path, audio, sr)
        print(f"Saved wav: {os.path.abspath(wav_path)}", flush=True)

    frames = analyze_audio(audio, sr, fmin=args.fmin, fmax=args.fmax)
    summary = summarize(frames, label=label)
    print_summary(summary)
    log_path = prefix if prefix.lower().endswith(".csv") else prefix + ".csv"
    # if prefix was .wav, still write .csv alongside
    if log_path.lower().endswith(".wav.csv"):
        log_path = log_path[:-8] + ".csv"
    elif prefix.lower().endswith(".wav"):
        log_path = prefix[:-4] + ".csv"
    write_frame_log(log_path, frames, summary)
    print(f"Saved log: {os.path.abspath(log_path)}", flush=True)
    print(
        "Tip: record I/II then III+ separately, then:\n"
        "  python tools/mic_pitch_analyze.py --compare mic_i.wav mic_iii.wav",
        flush=True,
    )


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nStopped.", flush=True)
        sys.exit(130)