#!/usr/bin/env python3
"""Through-air calibration for the Morse Code app (laptop speaker + mic).

Plays the app's impulse wire-code as percussive clicks out the laptop speaker
while recording the laptop mic, then runs the SAME transient-onset detector the
browser app uses (ported from marionette/tests/audio_analysis.py) and decodes
the result. Reports detection accuracy and recommends a `unitMs` so the app's
defaults can be tuned to this room/hardware — no robot required.

This MAKES NOISE. Turn the volume up and run it only when ready.

Run with the marionette venv (has sounddevice/scipy):
    /Users/remi/reachy_mini_apps/marionette/.venv/bin/python \
        tools/calibrate_audio.py --phrase "SOS HELLO" --sweep

Then set the winning unit in the app: lib/timing.js SPEED_PRESETS / Settings.
"""
from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

import numpy as np
import sounddevice as sd

# Reuse the exact detector the JS app was ported from.
MARIONETTE_TESTS = Path("/Users/remi/reachy_mini_apps/marionette/tests")
sys.path.insert(0, str(MARIONETTE_TESTS))
from audio_analysis import detect_transient_onsets  # noqa: E402

SR = 48000

# ── Morse + wire code (mirror of lib/morse.js, lib/wire.js, lib/timing.js) ──
CHAR_TO_MORSE = {
    "A": ".-", "B": "-...", "C": "-.-.", "D": "-..", "E": ".", "F": "..-.",
    "G": "--.", "H": "....", "I": "..", "J": ".---", "K": "-.-", "L": ".-..",
    "M": "--", "N": "-.", "O": "---", "P": ".--.", "Q": "--.-", "R": ".-.",
    "S": "...", "T": "-", "U": "..-", "V": "...-", "W": ".--", "X": "-..-",
    "Y": "-.--", "Z": "--..",
    "0": "-----", "1": ".----", "2": "..---", "3": "...--", "4": "....-",
    "5": ".....", "6": "-....", "7": "--...", "8": "---..", "9": "----.",
}
MORSE_TO_CHAR = {v: k for k, v in CHAR_TO_MORSE.items()}


def timing(unit_ms: int) -> dict:
    u = max(40, int(unit_ms))
    dah, elem, letter, word = 2 * u, 4 * u, 8 * u, 14 * u
    return {
        "unit": u, "dah": dah, "elem": elem, "letter": letter, "word": word,
        "click": min(60, int(u * 0.5)),
        "d_elem": round((dah * elem) ** 0.5),
        "d_letter": round((elem * letter) ** 0.5),
        "d_word": round((letter * word) ** 0.5),
    }


def text_to_onsets(text: str, t: dict) -> list[float]:
    onsets, cursor = [], 0.0
    words = [w for w in text.upper().split() if w]
    for wi, word in enumerate(words):
        if wi > 0:
            cursor += t["word"]
        letters = [CHAR_TO_MORSE[c] for c in word if c in CHAR_TO_MORSE]
        for li, code in enumerate(letters):
            if li > 0:
                cursor += t["letter"]
            for si, sym in enumerate(code):
                if si > 0:
                    cursor += t["elem"]
                onsets.append(cursor)
                if sym == "-":
                    cursor += t["dah"]
                    onsets.append(cursor)
    return onsets


def onsets_to_text(onsets_ms: list[float], t: dict) -> str:
    onsets = sorted(onsets_ms)
    if not onsets:
        return ""
    clusters, size, prev_last = [], 1, None
    i = 1
    while i <= len(onsets):
        ioi = onsets[i] - onsets[i - 1] if i < len(onsets) else float("inf")
        if i < len(onsets) and ioi < t["d_elem"]:
            size += 1
            i += 1
            continue
        gap = None if prev_last is None else onsets[i - size] - prev_last
        clusters.append((size, gap))
        prev_last = onsets[i - 1]
        size = 1
        i += 1
    morse_words, cur = [], ""
    for idx, (sz, gap) in enumerate(clusters):
        if idx > 0:
            if gap >= t["d_word"]:
                morse_words.append(cur)
                morse_words.append("/")
                cur = ""
            elif gap >= t["d_letter"]:
                cur += " "
        cur += "-" if sz >= 2 else "."
    morse_words.append(cur)
    out_words, letters = [], []
    for tok in " ".join(morse_words).split():
        if tok == "/":
            out_words.append("".join(letters))
            letters = []
        else:
            letters.append(MORSE_TO_CHAR.get(tok, "?"))
    out_words.append("".join(letters))
    return " ".join(w for w in out_words if w)


def render_clicks(onsets_ms: list[float], lead_ms=300, tail_ms=600, freq=2200.0) -> np.ndarray:
    dur_ms = lead_ms + (onsets_ms[-1] if onsets_ms else 0) + tail_ms
    n = int(SR * dur_ms / 1000)
    buf = np.zeros(n, dtype=np.float32)
    click_len = int(0.05 * SR)
    tt = np.arange(click_len) / SR
    env = np.exp(-tt / 0.012)
    blip = (0.7 * env * np.sin(2 * np.pi * freq * tt)).astype(np.float32)
    for t0 in onsets_ms:
        start = int((t0 + lead_ms) / 1000 * SR)
        end = min(start + click_len, n)
        buf[start:end] += blip[: end - start]
    return buf


def run_once(text: str, unit_ms: int, hp: float, ratio: float) -> dict:
    t = timing(unit_ms)
    onsets = text_to_onsets(text, t)
    audio = render_clicks(onsets, freq=2200.0)
    rec = sd.playrec(audio, samplerate=SR, channels=1, dtype="float32")
    sd.wait()
    rec = rec.flatten()
    detected = detect_transient_onsets(rec, SR, highpass_freq=hp, threshold_db=20 * np.log10(ratio))
    detected_ms = [d * 1000.0 for d in detected]
    decoded = onsets_to_text(detected_ms, t)
    return {
        "unit": unit_ms, "expected": len(onsets), "detected": len(detected_ms),
        "sent": text.upper(), "decoded": decoded, "ok": decoded == text.upper(),
    }


def main() -> int:
    ap = argparse.ArgumentParser(description="Morse audio through-air calibration")
    ap.add_argument("--phrase", default="SOS HELLO")
    ap.add_argument("--unit", type=int, default=120, help="base unit (ms)")
    ap.add_argument("--highpass", type=float, default=2000.0)
    ap.add_argument("--ratio", type=float, default=0.1, help="flux threshold ratio")
    ap.add_argument("--sweep", action="store_true", help="try several units, recommend one")
    args = ap.parse_args()

    print(f"Devices:\n{sd.query_devices()}\n")
    print(f"Phrase: {args.phrase!r}  (turn the volume up — this plays sound)\n")
    time.sleep(0.5)

    units = [160, 120, 90, 70] if args.sweep else [args.unit]
    results = []
    for u in units:
        print(f"— unit={u}ms …", flush=True)
        r = run_once(args.phrase, u, args.highpass, args.ratio)
        flag = "OK " if r["ok"] else "FAIL"
        print(f"  [{flag}] expected {r['expected']} hits, detected {r['detected']} | "
              f"decoded: {r['decoded']!r}")
        results.append(r)

    good = [r for r in results if r["ok"]]
    print("\n" + "=" * 60)
    if good:
        best = min(good, key=lambda r: r["unit"])  # fastest that still decodes
        print(f"RECOMMEND unitMs = {best['unit']}  (fastest reliable in this room)")
        print(f"  → set SPEED_PRESETS.normal / Settings unit to {best['unit']}")
    else:
        print("No unit decoded cleanly. Try: louder volume, slower (--unit 200),")
        print("lower --highpass (e.g. 1200), or higher --ratio (e.g. 0.15).")
    print("=" * 60)
    return 0 if good else 1


if __name__ == "__main__":
    sys.exit(main())