#!/usr/bin/env python3
"""
make_template.py - build the door-buzzer spectral template from an existing recording.

Usage:
  python tools/make_template.py --buzz buzzer.wav
         [--segments 1.2-2.0,5.5-6.4]        time ranges (s) that contain the buzzer
         [--background bg1.wav bg2.wav ...]  recordings WITHOUT the buzzer
         [--out main/template.h]             (main/test_vector.h is written next to it)
         [--gate-mult 4.0]                   auto-detect gate: noise floor * this
         [--min-rms 200] [--window-n 15] [--window-m 10] [--cooldown 0.5]
         [--threshold T]                     override threshold used in the false-trigger estimate

What it does:
  * loads any WAV (any rate / bit depth / channels), mixes to mono, resamples to FS_HZ
  * parses main/feature_params.h so the feature pipeline is exactly the firmware's
  * selects buzzer frames (given segments, or auto: RMS > 10th-percentile floor * gate-mult)
  * template = mean fingerprint of those frames, L2-normalised
  * with --background: scores every background frame, recommends a threshold and
    estimates false triggers / hour using the firmware's N-of-M window logic
  * writes main/template.h and main/test_vector.h (used by the firmware `selftest`)

Dependencies: numpy, scipy.
"""
import argparse
import math
import os
import re
import sys

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

MIN_FRAMES = 20


def die(msg):
    print("ERROR: " + msg, file=sys.stderr)
    sys.exit(1)


# ---------------------------------------------------------------- parameters
def parse_params(path):
    """Regex over '#define NAME value' lines of feature_params.h."""
    try:
        text = open(path, "r", encoding="utf-8").read()
    except OSError as e:
        die("cannot read %s: %s" % (path, e))
    raw = {}
    for m in re.finditer(r"^[ \t]*#define[ \t]+(\w+)[ \t]+([^\s/]+)", text, re.M):
        raw[m.group(1)] = m.group(2)
    need = ["FS_HZ", "FRAME_LEN", "HOP_LEN", "NUM_BANDS", "F_MIN_HZ", "F_MAX_HZ", "LOG_EPS", "FP_DETREND"]
    missing = [k for k in need if k not in raw]
    if missing:
        die("%s is missing #define(s): %s" % (path, ", ".join(missing)))
    P = {}
    for k in need:
        s = re.sub(r"[fFuUlL]+$", "", raw[k].strip("()"))   # 1e-10f -> 1e-10
        try:
            v = float(s)
        except ValueError:
            die("cannot parse value '%s' of %s in %s" % (raw[k], k, path))
        P[k] = v if k == "LOG_EPS" else int(round(v))
    if P["FRAME_LEN"] & (P["FRAME_LEN"] - 1):
        die("FRAME_LEN must be a power of two")
    return P


def band_edges(P):
    """Identical to fingerprint.c: log-spaced, round-half-up, >=1 bin per band."""
    n = P["NUM_BANDS"]
    ratio = P["F_MAX_HZ"] / P["F_MIN_HZ"]
    edges = []
    for b in range(n + 1):
        f = P["F_MIN_HZ"] * ratio ** (b / n)
        k = int(math.floor(f * P["FRAME_LEN"] / P["FS_HZ"] + 0.5))
        if b > 0 and k < edges[-1] + 1:
            k = edges[-1] + 1
        edges.append(k)
    if edges[-1] > P["FRAME_LEN"] // 2:
        die("band edges exceed Nyquist (%d > %d); lower NUM_BANDS or raise F_MIN_HZ"
            % (edges[-1], P["FRAME_LEN"] // 2))
    return np.array(edges, dtype=np.int64)


def hann(n):
    """Periodic Hann, as in fingerprint.c."""
    return 0.5 - 0.5 * np.cos(2.0 * np.pi * np.arange(n) / n)


# ------------------------------------------------------------------ features
def fingerprints(frames_i16, P, edges, window):
    """frames_i16: (F, FRAME_LEN) int16 -> fp (F, NUM_BANDS), rms (F,), valid (F,).
    Steps (same order as firmware): remove DC -> RMS -> /32768 * Hann -> FFT ->
    mean power per band -> ln(. + eps) -> subtract mean -> (detrend) -> L2 normalise."""
    L, N = P["FRAME_LEN"], P["NUM_BANDS"]
    x = frames_i16.astype(np.float64)
    x -= x.mean(axis=1, keepdims=True)
    rms = np.sqrt((x * x).mean(axis=1))
    spec = np.fft.rfft(x / 32768.0 * window, axis=1)[:, : L // 2]
    power = spec.real ** 2 + spec.imag ** 2
    sums = np.add.reduceat(power[:, : edges[-1]], edges[:-1], axis=1)
    v = np.log(sums / np.diff(edges) + P["LOG_EPS"])
    v -= v.mean(axis=1, keepdims=True)
    if P["FP_DETREND"]:
        xc = np.arange(N) - 0.5 * (N - 1)
        slope = (v @ xc) / (xc @ xc)
        v -= slope[:, None] * xc
    norm = np.linalg.norm(v, axis=1)
    valid = norm > 1e-3
    fp = np.zeros_like(v)
    fp[valid] = v[valid] / norm[valid, None]
    return fp, rms, valid


def all_frames(pcm_i16, P, edges, window, chunk=2000):
    L, H = P["FRAME_LEN"], P["HOP_LEN"]
    if len(pcm_i16) < L:
        return np.zeros((0, P["NUM_BANDS"])), np.zeros(0), np.zeros(0, bool)
    view = np.lib.stride_tricks.sliding_window_view(pcm_i16, L)[::H]
    fps, rmss, vals = [], [], []
    for i in range(0, len(view), chunk):
        fp, rms, ok = fingerprints(view[i:i + chunk], P, edges, window)
        fps.append(fp); rmss.append(rms); vals.append(ok)
    return np.vstack(fps), np.concatenate(rmss), np.concatenate(vals)


# ------------------------------------------------------------------- audio io
def load_wav(path, fs_target):
    """Any WAV -> mono int16-valued array at fs_target (values rounded like the mic path)."""
    try:
        fs, data = wavfile.read(path)
    except Exception as e:  # scipy raises various errors for odd formats
        die("cannot read %s as WAV (%s). If it is 24-bit/compressed, convert to 16-bit PCM or float "
            "WAV first (e.g. ffmpeg -i in.ext -ac 1 out.wav)." % (path, e))
    k, size = data.dtype.kind, data.dtype.itemsize
    if k == "u":
        x = (data.astype(np.float64) - 2 ** (8 * size - 1)) / 2 ** (8 * size - 1)
    elif k == "i":
        x = data.astype(np.float64) / 2 ** (8 * size - 1)
    elif k == "f":
        x = data.astype(np.float64)
    else:
        die("%s: unsupported sample type %s" % (path, data.dtype))
    if x.ndim == 2:
        x = x.mean(axis=1)
    if x.size == 0:
        die("%s contains no samples" % path)
    if fs != fs_target:
        g = math.gcd(int(fs), int(fs_target))
        x = resample_poly(x, fs_target // g, int(fs) // g)
    return np.clip(np.round(x * 32768.0), -32768, 32767).astype(np.int16)


# ----------------------------------------------------------------- selection
def parse_segments(s):
    out = []
    for tok in s.split(","):
        try:
            a, b = tok.strip().split("-")
            a, b = float(a), float(b)
        except ValueError:
            die("bad segment '%s' (expected START-END in seconds, e.g. 1.2-2.0)" % tok)
        if b <= a:
            die("segment '%s': end must be after start" % tok)
        out.append((a, b))
    return out


def ranges_text(idx, P):
    """Contiguous runs of frame indices -> 'start-end s' strings."""
    if len(idx) == 0:
        return []
    runs, start, prev = [], idx[0], idx[0]
    for i in idx[1:]:
        if i != prev + 1:
            runs.append((start, prev)); start = i
        prev = i
    runs.append((start, prev))
    H, L, FS = P["HOP_LEN"], P["FRAME_LEN"], P["FS_HZ"]
    return ["%.2f-%.2f s" % (a * H / FS, (b * H + L) / FS) for a, b in runs]


def select_frames(rms, valid, segments, gate_mult, P):
    H, L, FS = P["HOP_LEN"], P["FRAME_LEN"], P["FS_HZ"]
    n = len(rms)
    if segments:
        sel = np.zeros(n, bool)
        for a, b in segments:
            for i in range(n):
                if i * H >= a * FS and i * H + L <= b * FS:
                    sel[i] = True
        how = "segments"
    else:
        floor = max(np.percentile(rms, 10), 1.0)
        sel = rms > floor * gate_mult
        how = "auto (noise floor %.1f x %.1f = RMS gate %.1f)" % (floor, gate_mult, floor * gate_mult)
    sel &= valid
    return np.nonzero(sel)[0], how


# -------------------------------------------------------------------- report
def stats_line(label, a):
    return "%s: n=%d min=%.3f p5=%.3f p50=%.3f mean=%.3f" % (
        label, len(a), a.min(), np.percentile(a, 5), np.percentile(a, 50), a.mean())


def bar_chart(tpl, P, edges):
    scale = max(abs(tpl).max(), 1e-9)
    half = 24
    print("Template (band, centre Hz, value):")
    for b, v in enumerate(tpl):
        fc = P["FS_HZ"] * 0.5 * (edges[b] + edges[b + 1]) / P["FRAME_LEN"]
        n = int(round(abs(v) / scale * half))
        left = (" " * (half - n) + "#" * n) if v < 0 else " " * half
        right = ("#" * n) if v >= 0 else ""
        print("  %2d %6.0f Hz %+.3f %s|%s" % (b, fc, v, left, right))


def simulate_triggers(match, nframes_quiet_tail, N, M):
    """Simulate N-of-M detection and group matches inside the quiet tail."""
    hist, state, quiet_left, trig = 0, "IDLE", 0, 0
    mask = (1 << N) - 1
    for m in match:
        if state != "IDLE":
            if quiet_left <= 1:
                state, hist = "IDLE", 0
            else:
                quiet_left = nframes_quiet_tail if m else quiet_left - 1
                continue
        if state == "IDLE":
            hist = ((hist << 1) | int(m)) & 0xFFFFFFFF
            if bin(hist & mask).count("1") >= M:
                trig += 1
                state = "GROUPING"
                quiet_left = nframes_quiet_tail
                hist = 0
    return trig


# ------------------------------------------------------------------- writers
def fmt_floats(a, per_line=4):
    items = ["%.9ef" % v for v in a]
    lines = []
    for i in range(0, len(items), per_line):
        lines.append("    " + ", ".join(items[i:i + per_line]) + ",")
    return "\n".join(lines)


def write_template_h(path, tpl, count, thr, src, P):
    with open(path, "w") as f:
        f.write("// GENERATED by tools/make_template.py -- do not edit.\n")
        f.write("// Source recording: %s\n" % src)
        f.write("// Parameters: FS_HZ=%d FRAME_LEN=%d HOP_LEN=%d NUM_BANDS=%d F_MIN_HZ=%d F_MAX_HZ=%d "
                "LOG_EPS=%g FP_DETREND=%d\n" % (P["FS_HZ"], P["FRAME_LEN"], P["HOP_LEN"], P["NUM_BANDS"],
                                                 P["F_MIN_HZ"], P["F_MAX_HZ"], P["LOG_EPS"], P["FP_DETREND"]))
        f.write("#ifndef TEMPLATE_H\n#define TEMPLATE_H\n#include \"feature_params.h\"\n")
        f.write("#define TEMPLATE_VALID 1\n#define TEMPLATE_NUM_BANDS %d\n" % P["NUM_BANDS"])
        f.write("#define TEMPLATE_FRAME_COUNT %d\n" % count)
        if thr is not None:
            f.write("#define TEMPLATE_HAS_SUGGESTED_THRESHOLD 1\n#define TEMPLATE_SUGGESTED_THRESHOLD %.4ff\n" % thr)
        else:
            f.write("#define TEMPLATE_HAS_SUGGESTED_THRESHOLD 0\n")
        f.write("static const float TEMPLATE_FP[NUM_BANDS] = {\n%s\n};\n#endif\n" % fmt_floats(tpl))


def write_test_vector_h(path, pcm, fp, rms, frame_idx):
    with open(path, "w") as f:
        f.write("// GENERATED by tools/make_template.py -- do not edit.\n")
        f.write("// One int16 frame from the middle of the buzzer (frame #%d) and the Python-computed\n" % frame_idx)
        f.write("// fingerprint/RMS the firmware `selftest` must reproduce within 1e-3.\n")
        f.write("#ifndef TEST_VECTOR_H\n#define TEST_VECTOR_H\n#include <stdint.h>\n#include \"feature_params.h\"\n")
        f.write("#define TEST_VECTOR_VALID 1\n#define TEST_VECTOR_RMS %.6ff\n" % rms)
        f.write("static const int16_t TEST_VECTOR_PCM[FRAME_LEN] = {\n")
        for i in range(0, len(pcm), 16):
            f.write("    " + ", ".join(str(int(v)) for v in pcm[i:i + 16]) + ",\n")
        f.write("};\nstatic const float TEST_VECTOR_FP[NUM_BANDS] = {\n%s\n};\n#endif\n" % fmt_floats(fp))


# ---------------------------------------------------------------------- main
def main():
    here = os.path.dirname(os.path.abspath(__file__))
    root = os.path.dirname(here)
    ap = argparse.ArgumentParser(description="Build the buzzer spectral template (see file docstring).")
    ap.add_argument("--buzz", required=True, help="WAV containing the buzzer")
    ap.add_argument("--segments", help="comma list of START-END seconds, e.g. 1.2-2.0,5.5-6.4")
    ap.add_argument("--background", nargs="*", default=[], help="WAVs without the buzzer")
    ap.add_argument("--out", default=os.path.join(root, "main", "template.h"))
    ap.add_argument("--params", default=os.path.join(root, "main", "feature_params.h"))
    ap.add_argument("--gate-mult", type=float, default=4.0)
    ap.add_argument("--min-rms", type=float, default=200.0, help="firmware loudness gate (default Kconfig 200)")
    ap.add_argument("--window-n", type=int, default=15)
    ap.add_argument("--window-m", type=int, default=10)
    ap.add_argument("--cooldown", type=float, default=0.5,
                    help="quiet tail for grouping buzz frames, in seconds (default 0.5)")
    ap.add_argument("--threshold", type=float, help="threshold for the false-trigger estimate")
    a = ap.parse_args()

    P = parse_params(a.params)
    edges = band_edges(P)
    win = hann(P["FRAME_LEN"])
    print("Parameters from %s: %s" % (a.params, {k: v for k, v in P.items()}))
    print("Band edges (FFT bins):", edges.tolist())

    pcm = load_wav(a.buzz, P["FS_HZ"])
    print("Loaded %s: %.2f s at %d Hz" % (a.buzz, len(pcm) / P["FS_HZ"], P["FS_HZ"]))
    fps, rms, valid = all_frames(pcm, P, edges, win)
    if len(fps) == 0:
        die("recording shorter than one frame (%d samples)" % P["FRAME_LEN"])

    segs = parse_segments(a.segments) if a.segments else None
    idx, how = select_frames(rms, valid, segs, a.gate_mult, P)
    print("Selected %d of %d frames via %s" % (len(idx), len(fps), how))
    print("Time ranges:", ", ".join(ranges_text(idx, P)) or "(none)")
    if len(idx) < MIN_FRAMES:
        die("only %d buzzer frames selected, need >= %d. Give --segments, lower --gate-mult, "
            "or use a longer/louder recording." % (len(idx), MIN_FRAMES))

    tpl = fps[idx].mean(axis=0)
    nrm = np.linalg.norm(tpl)
    if nrm < 1e-3:
        die("mean fingerprint is degenerate (flat spectrum)")
    tpl /= nrm

    buzz_scores = fps[idx] @ tpl
    print("\n" + stats_line("Buzzer frames vs template", buzz_scores))
    p5 = float(np.percentile(buzz_scores, 5))
    print()
    bar_chart(tpl, P, edges)

    thr_suggest = None
    if a.background:
        bg_all, n_hours = [], 0.0
        sims = []
        thr_for_sim = a.threshold
        per_file = []
        for path in a.background:
            bpcm = load_wav(path, P["FS_HZ"])
            bfp, brms, bval = all_frames(bpcm, P, edges, win)
            if len(bfp) == 0:
                print("WARNING: %s shorter than a frame, skipped" % path)
                continue
            sc = bfp @ tpl
            sc[~bval] = 0.0
            per_file.append((path, sc, brms, len(bpcm) / P["FS_HZ"]))
            bg_all.append(sc)
        if not bg_all:
            die("no usable background audio")
        bg = np.concatenate(bg_all)
        p99 = float(np.percentile(bg, 99))
        print("\nBackground (%d frames from %d file(s)): min=%.3f mean=%.3f p95=%.3f p99=%.3f max=%.3f" % (
            len(bg), len(per_file), bg.min(), bg.mean(), np.percentile(bg, 95), p99, bg.max()))
        thr_suggest = float(np.clip(0.5 * (p99 + p5), 0.50, 0.98))
        print("Recommended threshold: %.3f  (midpoint of background p99 %.3f and buzzer p5 %.3f)" % (thr_suggest, p99, p5))
        if p5 <= p99:
            print("!" * 70)
            print("!! WARNING: background p99 (%.3f) >= buzzer p5 (%.3f): DISTRIBUTIONS OVERLAP." % (p99, p5))
            print("!! Expect false alarms or missed buzzes. Use a cleaner/louder buzzer recording,")
            print("!! more specific --segments, or try FP_DETREND 0 (then rerun this tool).")
            print("!" * 70)
        thr_use = a.threshold if a.threshold is not None else thr_suggest
        cd_frames = int(round(a.cooldown * P["FS_HZ"] / P["HOP_LEN"]))
        trig, secs = 0, 0.0
        for path, sc, brms, dur in per_file:
            match = (brms >= a.min_rms) & (sc >= thr_use)
            trig += simulate_triggers(match, cd_frames, a.window_n, a.window_m)
            secs += dur
        print("Frames passing the loudness gate (RMS>=%.0f): %.1f%%" % (
            a.min_rms, 100.0 * np.mean(np.concatenate([f[2] for f in per_file]) >= a.min_rms)))
        print("False-trigger estimate at threshold %.3f with %d-of-%d window, %.0f ms quiet tail: "
              "%d buzz event(s) in %.1f min = %.2f per hour" % (
                  thr_use, a.window_m, a.window_n, a.cooldown * 1000.0,
                  trig, secs / 60.0, trig / (secs / 3600.0)))

    out_dir = os.path.dirname(os.path.abspath(a.out)) or "."
    os.makedirs(out_dir, exist_ok=True)
    write_template_h(a.out, tpl, len(idx), thr_suggest, os.path.basename(a.buzz), P)
    tv_path = os.path.join(out_dir, "test_vector.h")
    mid = int(idx[len(idx) // 2])
    frame = pcm[mid * P["HOP_LEN"]: mid * P["HOP_LEN"] + P["FRAME_LEN"]]
    tfp, trms, tok = fingerprints(frame[None, :], P, edges, win)
    if not tok[0]:
        die("test-vector frame has a degenerate fingerprint")
    write_test_vector_h(tv_path, frame, tfp[0], float(trms[0]), mid)
    print("\nWrote %s\nWrote %s\nNext: idf.py build flash monitor, then `selftest`, then `selfcheck`." % (a.out, tv_path))


if __name__ == "__main__":
    main()