#!/usr/bin/env python3
"""Find where each ORIGINAL-GAME sample sits inside the record's vocal hooks.

The Mortal Kombat hooks were separated out of the record by UVR, and they sound
it: the stem has no energy below ~400 Hz, so "Test your might" arrives thin and
the whole SUB_MIX / SUB2_MIX / PRESENCE apparatus in mk-mix-hooks.mjs exists to
fake the chest back in. But the record did not invent those lines -- every one
except "MORTAL KOMBAT!" is a sample of the 1992 arcade announcer, and the arcade
samples are clean and full-bodied. So the hooks can be REBUILT from the game's
own audio, placed exactly where the record plays them.

This measures the two things that takes:

  rate     the record plays every arcade sample ONE SEMITONE DOWN. Measured by
           raw-waveform correlation over rates 0.930-0.960 at 0.001 steps: the
           best-correlated takes all peak at 0.944 (r 0.63-0.71), i.e.
           2^(-1/12) = 0.94387. One global rate, not one per line -- a 1993
           sampler played the whole bank at one pitch.

  onset    where each line starts, in RECORD seconds, per occurrence. Band
           energy envelopes (300-3500 Hz, 2 ms hop) find it to a frame; then
           raw-waveform correlation within +-60 ms pins it to a sample -- but
           only where that correlation is real (r >= 0.3). "Finish him!",
           "Excellent!" and "Fight!" correlate at 0.08-0.23 raw (the record
           treats them more heavily), so they keep the envelope answer.

A hook holding two lines ("Excellent! Kano!") searches for the second only
after the first is mostly over: unconstrained, hook-18's "Kano" matched 0.29 s
into its own "Excellent", which is not a place a 0.7 s name can start.

Input is what exists on a machine without the song data: the hooks UI's
catalogue (mk-hooks.data.js) and its padded stem cuts (mk-hooks.audio/*.mp3,
each opening `pad` seconds before the catalogue window). The game pack is the
MKWarehouse MK1 set; see GAME_LINES.

  align-game-hooks.py <reports-dir> <game-pack-dir> <out.json>

Writes the tracked mapping build-game-hooks.mjs consumes. Nothing else.
"""
import json
import os
import re
import subprocess
import sys

import numpy as np

SR = 48000
HOP = 96  # 2 ms
SEMITONE_DOWN = 2 ** (-1 / 12)

# Line text (as the catalogue spells it) -> arcade sample. The catalogue's
# spellings are the separator-era transcription ("Senya!", "MORDO GRANDE!").
GAME_LINES = {
    "Test your might.": "mk1-00381",
    "Finish him!": "mk1-00378",
    "Excellent!": "mk1-00377",
    "Fight!": "mk1-00368",
    "Kano!": "mk1-00409",
    "Liu Kang!": "mk1-00410",
    "Raiden!": "mk1-00411",
    "Johnny Cage!": "mk1-00412",
    "Scorpion!": "mk1-00413",
    "Sub Zero!": "mk1-00414",
    "Senya!": "mk1-00415",
}
# Lines with no arcade source. "MORTAL KOMBAT!" is from the 1993 home-release TV
# spot, not the game; "I'm gonna kill you!" is not in the MK1 announcer bank.
KEEP_LINES = {"MORDO GRANDE!", "I'm gonna kill you!"}
RAW_MIN = 0.3


def decode(path):
    raw = subprocess.run(
        ["ffmpeg", "-v", "error", "-i", path, "-ac", "1", "-ar", str(SR), "-f", "s16le", "-"],
        check=True, capture_output=True).stdout
    return np.frombuffer(raw, dtype=np.int16).astype(np.float64) / 32768


def resample(x, rate):
    """Play x at `rate` (a sampler's pitch-and-speed change together)."""
    n = int(len(x) / rate)
    return np.interp(np.arange(n) * rate, np.arange(len(x)), x)


def bandpass(x, lo=150, hi=3800):
    X = np.fft.rfft(x)
    f = np.fft.rfftfreq(len(x), 1 / SR)
    X[(f < lo) | (f > hi)] = 0
    return np.fft.irfft(X, len(x))


def envelope(x, lo=300, hi=3500):
    N = 1024
    win = np.hanning(N)
    k0, k1 = int(lo / (SR / N)), int(hi / (SR / N))
    out = []
    for off in range(0, max(1, len(x) - N), HOP):
        s = x[off:off + N]
        if len(s) < N:
            s = np.pad(s, (0, N - len(s)))
        out.append(np.log((np.abs(np.fft.rfft(s * win))[k0:k1] ** 2).sum() + 1e-9))
    return np.array(out)


def env_search(hay, needle, lo_s, hi_s):
    """Best needle START (seconds) in [lo_s, hi_s], by normalised envelope correlation."""
    L = len(needle)
    nn = (needle - needle.mean()) / (needle.std() + 1e-9)
    padded = np.pad(hay, (L, L), constant_values=hay.min())
    best, best_i, scores = -9.0, 0, []
    for i in range(len(padded) - L + 1):
        t = (i - L) * HOP / SR
        if t < lo_s or t > hi_s:
            continue
        s = padded[i:i + L]
        sd = s.std()
        v = -9.0 if sd < 1e-6 else float(np.dot((s - s.mean()) / sd, nn) / L)
        scores.append(v)
        if v > best:
            best, best_i = v, i
    sc = np.array([v for v in scores if v > -9])
    z = (best - sc.mean()) / (sc.std() + 1e-9) if len(sc) > 1 else 0.0
    return (best_i - L) * HOP / SR, best, z


def raw_search(hay, needle, around_s, span_s=0.06):
    """Sample-accurate start near around_s, with the normalised correlation there."""
    n, m = len(needle), len(hay) + 1  # m candidate starts, 0..len(hay)
    h = np.pad(hay, (0, n))
    N = 1 << int(np.ceil(np.log2(len(h) + n)))
    c = np.fft.irfft(np.fft.rfft(h, N) * np.conj(np.fft.rfft(needle, N)), N)[:m]
    cs = np.concatenate([[0], np.cumsum(h * h)])
    e = np.sqrt(np.maximum(cs[n:n + m] - cs[:m], 0) + 1e-12)
    c = c / (e * np.sqrt((needle * needle).sum()) + 1e-12)
    lo = max(0, int((around_s - span_s) * SR))
    hi = min(len(c), int((around_s + span_s) * SR))
    if hi <= lo:
        return around_s, 0.0
    i = lo + int(np.argmax(c[lo:hi]))
    return i / SR, float(c[i])


def measure_rate(pairs):
    """Median best rate over the takes that correlate for real."""
    found = []
    for hay, needle, around in pairs:
        best = (-1.0, 1.0)
        for r in np.arange(0.930, 0.9601, 0.001):
            _, v = raw_search(hay, resample(needle, r), around, 0.05)
            if v > best[0]:
                best = (v, float(r))
        if best[0] >= RAW_MIN:
            found.append(best)
    rates = sorted(r for _, r in found)
    return (rates[len(rates) // 2] if rates else None), found


def main():
    reports, pack, out = sys.argv[1:4]
    js = open(os.path.join(reports, "mk-hooks.data.js")).read()
    cat = json.loads(js[js.index("{"):js.rindex("}") + 1])
    pad = cat.get("pad", 1)
    clips = cat["clips"]
    game = {s: bandpass(decode(os.path.join(pack, s + ".mp3"))) for s in set(GAME_LINES.values())}
    game_len = {s: len(x) / SR for s, x in game.items()}

    hooks = {}
    for c in clips:
        hooks[c["id"]] = bandpass(decode(os.path.join(reports, "mk-hooks.audio", c["id"] + ".mp3")))

    def parts_of(text):
        return [p.strip() for p in re.split(r"(?<=[!.])\s+", text) if p.strip()]

    # Rate first, on the single-line name hooks (the cleanest correlations).
    probe = []
    for c in clips:
        ps = parts_of(c["text"])
        if len(ps) == 1 and ps[0] in GAME_LINES and ps[0] not in ("Finish him!", "Excellent!", "Fight!"):
            s = GAME_LINES[ps[0]]
            he = envelope(hooks[c["id"]])
            at, _, _ = env_search(he, envelope(resample(game[s], SEMITONE_DOWN)), 0, len(hooks[c["id"]]) / SR)
            probe.append((hooks[c["id"]], game[s], at))
    measured, found = measure_rate(probe)
    print(f"rate: median {measured:.3f} over {len(found)} takes with raw r >= {RAW_MIN}"
          f" (one semitone down = {SEMITONE_DOWN:.5f})", flush=True)
    rate = SEMITONE_DOWN

    env_cache = {s: envelope(resample(x, rate)) for s, x in game.items()}
    rs_cache = {s: resample(x, rate) for s, x in game.items()}

    rows = []
    for c in clips:
        hay = hooks[c["id"]]
        he = envelope(hay)
        t0 = c["t0"]
        parts, after = [], 0.0
        for text in parts_of(c["text"]):
            if text in KEEP_LINES or text not in GAME_LINES:
                parts.append({"line": text, "keep": True})
                continue
            s = GAME_LINES[text]
            at_env, r_env, z = env_search(he, env_cache[s], after, len(hay) / SR)
            at_raw, r_raw = raw_search(hay, rs_cache[s], at_env)
            use_raw = r_raw >= RAW_MIN
            at = at_raw if use_raw else at_env
            parts.append({
                "line": text,
                "sample": s,
                "recOnset": round(t0 + at, 4),
                "by": "raw" if use_raw else "envelope",
                "rRaw": round(r_raw, 3),
                "rEnv": round(r_env, 3),
                "z": round(z, 1),
            })
            # The next line cannot begin until this one is most of the way out.
            after = at + 0.8 * game_len[s] / rate
        # A kept line that shares a hook with a game line ends where the game line
        # starts -- the kept audio is cut there.
        for i, p in enumerate(parts):
            if p.get("keep"):
                nxt = next((q for q in parts[i + 1:] if not q.get("keep")), None)
                if nxt:
                    p["keepUntil"] = round(nxt["recOnset"] - 0.03, 4)
        rows.append({
            "id": c["id"],
            "text": c["text"],
            "catStart": c["cat"][0],
            "catEnd": c["cat"][1],
            "clipStart": t0,
            "parts": parts,
        })
        print(c["id"], json.dumps(parts), flush=True)

    doc = {
        "about": "Original-game replacements for the Mortal Kombat vocal hooks. Written by "
                 "align-game-hooks.py; consumed by build-game-hooks.mjs.",
        "pack": {
            "name": "MKWarehouse Mortal Kombat (1992 arcade) sounds",
            "url": "https://www.mortalkombatwarehouse.com/mk1/sounds/mkw_mk1_sounds.zip",
            "dir": "mk1-samples/zip/announcer",
        },
        "rate": rate,
        "rateMeasured": measured,
        "rateTakes": len(found),
        "lines": GAME_LINES,
        "kept": sorted(KEEP_LINES),
        "catalogue": "mk-hooks.data.js",
        "pad": pad,
        "hooks": rows,
    }
    with open(out, "w") as f:
        json.dump(doc, f, indent=1)
        f.write("\n")
    print(f"wrote {out}")


if __name__ == "__main__":
    main()