"""Lists the simulated clips most similar to a chosen simulated clip under the
perceptual log-mel distance of match.py, and writes a listening page.

    python3 ml/similar.py QUERY_ID [--data dataset] [--top 5] [--controls 50,500,2500]

The log-mel features of every clip are cached in DATA/ml/mel.npz. The page
DATA/similar-QUERY_ID.html has the query, the top matches, and a few control
clips at given ranks, so the measure can be judged by ear against the ranks.
"""

import argparse
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from match import read_wav_any, segment, melspec  # noqa: E402


def mel_features(data):
    cache = os.path.join(data, "ml", "mel.npz")
    rows = {json.loads(l)["id"]: json.loads(l) for l in open(os.path.join(data, "index.jsonl")) if l.strip()}
    ids = sorted(i for i in rows if os.path.exists(os.path.join(data, "clips", i + ".wav")))
    if os.path.exists(cache):
        z = np.load(cache, allow_pickle=True)
        if list(z["ids"]) == ids:
            return rows, ids, z["mel"]
    os.makedirs(os.path.join(data, "ml"), exist_ok=True)
    mel = np.stack([melspec(segment(read_wav_any(os.path.join(data, "clips", i + ".wav")), 0)) for i in ids]).astype(np.float32)
    np.savez(cache, ids=np.array(ids), mel=mel)
    return rows, ids, mel


def cells(r):
    return (f"<td>{r['E'] / 1e9:.0f}</td><td>{r['rho']:.0f}</td><td>{r['h'] * 1e3:.2f}</td><td>{r['Lx'] / 0.0254:.1f}×{r['Ly'] / 0.0254:.1f}</td>"
            f"<td>{r['strikeVelocity']:.1f}</td><td>{r['strikeFracX']:.2f}, {r['strikeFracY']:.2f}</td><td>{r['sigma0']:.2f}</td><td>{r['sigma1'] * 1e4:.1f}</td>"
            f"<td>{r['peakToPeak'] * 1e3:.1f}</td><td class='l'><div class='feat' title=\"{r['features']}\">{r['features']}</div></td>")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("query")
    ap.add_argument("--data", default="dataset")
    ap.add_argument("--top", type=int, default=5)
    ap.add_argument("--controls", default="50,500,2500")
    args = ap.parse_args()
    rows, ids, mel = mel_features(args.data)
    q = ids.index(args.query)
    d = np.abs(mel - mel[q]).mean((1, 2))
    order = np.argsort(d)
    assert order[0] == q
    ranked = [(k, ids[i], d[i]) for k, i in enumerate(order)]
    sel = [("query", *ranked[0][1:])] + [(f"rank {k}", i, dd) for k, i, dd in ranked[1:args.top + 1]]
    for c in [int(x) for x in args.controls.split(",") if x]:
        if c < len(ranked):
            sel.append((f"control, rank {c}", *ranked[c][1:]))
    print(f"query {args.query}; median distance over all clips {np.median(d):.2f} dB")
    for label, i, dd in sel:
        r = rows[i]
        print(f"{label:>18} {i} {dd:5.2f} dB | E {r['E'] / 1e9:.0f} rho {r['rho']:.0f} h {r['h'] * 1e3:.2f} {r['Lx'] / 0.0254:.1f}x{r['Ly'] / 0.0254:.1f} in v {r['strikeVelocity']:.1f} sig0 {r['sigma0']:.2f} sig1 {r['sigma1'] * 1e4:.1f} p2p {r['peakToPeak'] * 1e3:.1f}")
    gal = open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "scripts", "gallery.html")).read()
    style = gal[gal.index("<style>") + 7:gal.index("</style>")]
    js = gal[gal.index("// Spectrogram:"):gal.index("$(\"reload\").onclick")]
    body = "".join(
        f"<tr><td class='l'>{label}</td><td class='l'>{i}</td><td>{dd:.2f}</td><td class='l'><audio controls preload='none' src='clips/{i}.wav'></audio></td>{cells(rows[i])}</tr>"
        f"<tr><td colspan='14' class='l'><canvas width='900' height='120' data-src='clips/{i}.wav'></canvas></td></tr>"
        for label, i, dd in sel)
    html = f"""<!doctype html><meta charset="utf-8"><title>clips similar to {args.query}</title><style>{style}</style>
<h1>Simulated clips most similar to {args.query}</h1>
<p>Distance: mean absolute dB difference between level-normalized log-mel spectrograms (48 bands, 60 Hz to 6 kHz, 100 ms frames, floored at −60 dB). Median distance over all {len(ids)} clips: {np.median(d):.2f} dB. The control rows are clips at the given ranks, for comparison by ear.</p>
<table><thead><tr><th class="l">role</th><th class="l">id</th><th>dist dB</th><th class="l">audio</th><th>E GPa</th><th>ρ</th><th>h mm</th><th>W×H in</th><th>v</th><th>strike</th><th>σ0</th><th>σ1·10⁴</th><th>p-p mm</th><th class="l">features</th></tr></thead><tbody>{body}</tbody></table>
<script>const $ = (id) => document.getElementById(id);
{js}
for (const c of document.querySelectorAll("canvas")) drawSpec(c, c.dataset.src);
</script>"""
    out = os.path.join(args.data, f"similar-{args.query}.html")
    open(out, "w").write(html)
    print("wrote", out)


if __name__ == "__main__":
    main()