#!/usr/bin/env python3
r"""set_finder — recover a set's tracklist from a single room recording.

THE PROBLEM
    A gig was captured on a handheld stage mic. There are no stems, no gig-log,
    and PLN's own verdict on the audio is "bad rec cause the whole soundsystem had
    effects and at some points ppl messed with it". Every boundary lens the
    toolbox already has needs something this recording does not have:

      orbit-activity  -> needs stems, there are none
      gig-log `track` -> the 1 Hz recorder never ran that night
      gap detection   -> a livecoded set does not stop between tracks

    What is left is the audio itself, and one large advantage: we know what was
    PREPPED, so this is a MATCHING problem over ~23 candidates rather than open
    discovery. (PLN: "the set was slightly altered" — so candidates, not truth.
    Absence of a candidate is weak evidence; presence is strong.)

TWO LENSES, AND WHY BOTH
    1. TEMPO segments the set but cannot name it. Five candidates sit at 120 BPM
       and three at 124, so tempo alone leaves ties it can never break. Tempo from
       audio is also only determined up to a factor of two, so it is reported as a
       family (see reference_tempo_from_audio).

    2. SAMPLE SIGNATURE names a segment. sample_tfidf.py already computes which
       sounds IDENTIFY a track — `take5` and `love_parade` are df=1, `jungle_breaks`
       is in 80 tracks and worthless. We fingerprint only the rare ones.

    Tempo says WHERE the seams are; the signature says WHICH track. Neither alone
    is enough, and saying which lens produced a given claim is the point.

WHY CONSTELLATION FINGERPRINTS AND NOT CORRELATION
    The reference is a dry one-shot from disk; the query is that sound through a
    PA, a room, a limiter, someone's fingers on the desk, and a 320 kbps MP3.
    Cross-correlation dies on all of that. Peak constellations survive it: they
    encode only WHICH time-frequency points stick up locally, which is preserved
    under EQ, level, and moderate reverb. Same reason it works on a phone in a bar.

HONEST LIMITS — read before trusting output
    * A detection means "this sound occurred here", NOT "this track played here".
      PLN reuses banks live; the tracklist claim is the ASSEMBLY step's, and it is
      reported with its evidence so a wrong call is visible rather than laundered.
    * Recall is not symmetric with precision here. A track whose signature never
      sounded (a gate left closed, a stem never played) is invisible. Silence is
      not absence of the track.
    * Vote counts are NOT comparable between banks: a 582-file bank offers far
      more chances to match than a 4-file one. Scores are normalised per bank, and
      the raw counts are printed so the asymmetry stays visible.

USAGE
    set_finder.py fingerprint            # build the reference DB from rare banks
    set_finder.py scan REC.mp3           # detections + tempo, writes JSON
    set_finder.py assemble REC.mp3       # the proposed tracklist, with evidence
"""

from __future__ import annotations

import argparse
import collections
import json
import pathlib
import re
import subprocess
import sys
import time

import numpy as np

REPO = pathlib.Path("/home/pln/Work/Sound/Tidal")
TFIDF = REPO / "armada" / "tide-table" / "sample_tfidf.json"
DIRT = pathlib.Path.home() / ".local/share/SuperCollider/downloaded-quarks/Dirt-Samples"
DB_OUT = REPO / "armada" / "tide-table" / "set_finder_db.npz"
SCAN_OUT = REPO / "armada" / "tide-table" / "set_finder_scan.json"

SR = 22050
N_FFT = 2048
HOP = 512
# A sound must be rare to be a fingerprint. df<=6 keeps `take5`(1) and
# `bogdan_grime`(5) while rejecting `jungle_breaks`(80) and `bassWarsaw`(266).
MAX_DF = 6
MAX_FILES_PER_BANK = 12  # bounds DB size; longest files first (most hashes)
PEAK_NEIGHBOURHOOD = (24, 12)  # (freq bins, time frames) for local-max filter
PEAKS_PER_FRAME_CAP = 5
FAN_OUT = 8  # peaks paired forward per anchor
TARGET_DT = (2, 40)  # frames; ~46ms..0.93s at hop 512 / 22050


def _lib():
    try:
        import librosa  # noqa: PLC0415
        return librosa
    except ImportError:
        sys.exit("set_finder: librosa is required (it is present in system python here)")


# ------------------------------------------------------------- fingerprinting ----


def peaks_of(spec_db: np.ndarray) -> np.ndarray:
    """(freq_bin, frame) local maxima, capped per frame. spec_db is [freq, time]."""
    from scipy.ndimage import maximum_filter  # noqa: PLC0415

    mx = maximum_filter(spec_db, size=PEAK_NEIGHBOURHOOD, mode="nearest")
    ispeak = (spec_db == mx) & (spec_db > spec_db.max() - 60.0)
    f_idx, t_idx = np.nonzero(ispeak)
    if len(t_idx) == 0:
        return np.empty((0, 2), dtype=np.int32)
    # cap per frame, keeping the loudest
    order = np.lexsort((-spec_db[f_idx, t_idx], t_idx))
    f_idx, t_idx = f_idx[order], t_idx[order]
    keep = np.ones(len(t_idx), dtype=bool)
    seen: dict[int, int] = {}
    for i, t in enumerate(t_idx):
        c = seen.get(t, 0)
        if c >= PEAKS_PER_FRAME_CAP:
            keep[i] = False
        else:
            seen[t] = c + 1
    return np.stack([f_idx[keep], t_idx[keep]], axis=1).astype(np.int32)


def hashes_of(pk: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Constellation pairs -> (hash, anchor_frame).

    hash packs (f1, f2, dt) into one int64. Quantising f to 10 bits (1024 bins is
    already the STFT resolution) and dt to 6 bits keeps it collision-cheap.
    """
    if len(pk) == 0:
        return np.empty(0, np.int64), np.empty(0, np.int32)
    H, A = [], []
    n = len(pk)
    for i in range(n):
        f1, t1 = pk[i]
        j = i + 1
        fan = 0
        while j < n and fan < FAN_OUT:
            f2, t2 = pk[j]
            dt = t2 - t1
            if dt > TARGET_DT[1]:
                break
            if dt >= TARGET_DT[0]:
                H.append((int(f1) << 20) | (int(f2) << 10) | int(min(dt, 63)))
                A.append(t1)
                fan += 1
            j += 1
    return np.array(H, np.int64), np.array(A, np.int32)


def fingerprint_file(path: pathlib.Path) -> tuple[np.ndarray, np.ndarray]:
    librosa = _lib()
    y, _ = librosa.load(str(path), sr=SR, mono=True)
    if len(y) < N_FFT:
        return np.empty(0, np.int64), np.empty(0, np.int32)
    S = np.abs(librosa.stft(y, n_fft=N_FFT, hop_length=HOP))
    return hashes_of(peaks_of(librosa.amplitude_to_db(S, ref=np.max)))


def candidate_tracks() -> list[str]:
    out = subprocess.run(
        ["git", "-C", str(REPO), "show", "--stat", "--oneline", "eb9221d"],
        capture_output=True, text=True,
    ).stdout
    return sorted(set(re.findall(r"live/[a-zA-Z0-9_/]*\.tidal", out)))


def rare_banks() -> dict[str, list[str]]:
    """track stem -> [rare bank names], from the authored tf-idf report."""
    if not TFIDF.exists():
        sys.exit(f"set_finder: {TFIDF} missing — run tools/sample_tfidf.py first")
    rep = json.load(open(TFIDF))
    by_stem = {pathlib.Path(k).stem: k for k in rep["tracks"] if k.startswith("live/")}
    res: dict[str, list[str]] = {}
    for c in candidate_tracks():
        stem = pathlib.Path(c).stem
        k = by_stem.get(stem)
        if not k:
            continue
        rare = [
            t["sample"] for t in rep["tracks"][k].get("top_tfidf", [])
            if t["df"] <= MAX_DF and (DIRT / t["sample"]).exists()
        ]
        if rare:
            res[stem] = rare
    return res


def cmd_fingerprint(args) -> int:
    banks = rare_banks()
    all_banks = sorted({b for v in banks.values() for b in v})
    print(f"{len(banks)} candidate tracks carry {len(all_banks)} rare banks (df<={MAX_DF})")
    H, B, A, F = [], [], [], []
    bank_ids = {b: i for i, b in enumerate(all_banks)}
    t0 = time.time()
    for b in all_banks:
        files = sorted(
            (p for p in (DIRT / b).iterdir() if p.suffix.lower() in (".wav", ".flac", ".aiff")),
            key=lambda p: -p.stat().st_size,
        )[:MAX_FILES_PER_BANK]
        got = 0
        for fi, p in enumerate(files):
            try:
                h, a = fingerprint_file(p)
            except Exception:
                continue
            if len(h) == 0:
                continue
            H.append(h); A.append(a)
            B.append(np.full(len(h), bank_ids[b], np.int16))
            F.append(np.full(len(h), fi, np.int16))
            got += len(h)
        print(f"  {b:<22} {len(files):>3} files  {got:>7} hashes")
    if not H:
        sys.exit("set_finder: no hashes built")
    np.savez_compressed(
        DB_OUT,
        hashes=np.concatenate(H), banks=np.concatenate(B),
        anchors=np.concatenate(A), fileidx=np.concatenate(F),
        bank_names=np.array(all_banks),
        track_banks=np.array(json.dumps(banks)),
    )
    total = sum(len(h) for h in H)
    print(f"\n{total} hashes -> {DB_OUT.name} "
          f"({DB_OUT.stat().st_size/1e6:.1f} MB) in {time.time()-t0:.0f}s")
    return 0


# --------------------------------------------------------------------- scan ----


def cmd_scan(args) -> int:
    librosa = _lib()
    if not DB_OUT.exists():
        sys.exit("set_finder: run `fingerprint` first")
    db = np.load(DB_OUT, allow_pickle=False)
    bank_names = [str(x) for x in db["bank_names"]]
    track_banks = json.loads(str(db["track_banks"]))
    # hash -> indices, built once
    order = np.argsort(db["hashes"], kind="stable")
    h_sorted = db["hashes"][order]
    b_sorted = db["banks"][order]
    a_sorted = db["anchors"][order]

    rec = pathlib.Path(args.recording)
    dur = float(subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "csv=p=0", str(rec)], capture_output=True, text=True).stdout.strip())
    print(f"scanning {rec.name}  {dur/60:.1f} min, chunks of {args.chunk}s\n")

    votes: dict[tuple[int, int], int] = collections.defaultdict(int)
    bank_hits: dict[int, list[float]] = collections.defaultdict(list)
    t0 = time.time()
    pos = 0.0
    while pos < dur:
        y, _ = librosa.load(str(rec), sr=SR, mono=True, offset=pos, duration=args.chunk)
        if len(y) < N_FFT:
            break
        S = np.abs(librosa.stft(y, n_fft=N_FFT, hop_length=HOP))
        pk = peaks_of(librosa.amplitude_to_db(S, ref=np.max))
        qh, qa = hashes_of(pk)
        if len(qh):
            li = np.searchsorted(h_sorted, qh, side="left")
            ri = np.searchsorted(h_sorted, qh, side="right")
            for k in range(len(qh)):
                lo, hi = li[k], ri[k]
                if hi <= lo or hi - lo > args.max_ambiguity:
                    continue
                for m in range(lo, hi):
                    bank = int(b_sorted[m])
                    off = int(qa[k]) - int(a_sorted[m])
                    votes[(bank, off // args.offset_bin)] += 1
                    bank_hits[bank].append(pos + qa[k] * HOP / SR)
        pct = min(100.0, 100.0 * (pos + args.chunk) / dur)
        print(f"\r  {pct:5.1f}%  {len(votes)} vote bins", end="", flush=True)
        pos += args.chunk
    print(f"\n  scanned in {time.time()-t0:.0f}s\n")

    # tempo curve over the whole file (1-D, cheap)
    print("measuring tempo …")
    y, _ = librosa.load(str(rec), sr=SR, mono=True)
    onset = librosa.onset.onset_strength(y=y, sr=SR, hop_length=HOP)
    win = int(args.tempo_window * SR / HOP)
    tempi = []
    for s in range(0, len(onset) - win, win // 2):
        t = librosa.feature.tempo(onset_envelope=onset[s:s + win], sr=SR, hop_length=HOP)
        tempi.append((s * HOP / SR, float(t[0])))
    print(f"  {len(tempi)} tempo windows of {args.tempo_window}s\n")

    result = {
        "recording": str(rec), "duration_s": dur,
        "bank_detections": {
            bank_names[b]: sorted(set(round(t, 1) for t in ts))
            for b, ts in bank_hits.items()
        },
        "bank_hit_counts": {bank_names[b]: len(ts) for b, ts in bank_hits.items()},
        "tempo_curve": tempi,
        "track_banks": track_banks,
        "params": {"max_df": MAX_DF, "max_files_per_bank": MAX_FILES_PER_BANK,
                   "sr": SR, "hop": HOP, "peaks_per_frame": PEAKS_PER_FRAME_CAP},
    }
    SCAN_OUT.write_text(json.dumps(result, indent=1))
    print(f"wrote {SCAN_OUT.name}")

    print("\nbank hit counts (RAW — not comparable between banks, see the docstring):")
    for b, n in sorted(result["bank_hit_counts"].items(), key=lambda kv: -kv[1])[:20]:
        print(f"  {b:<22} {n:>7}")
    return 0


# ------------------------------------------------------------------- locate ----
#
# `scan` counted raw hash collisions, and that is not the algorithm. Three fixes,
# each traceable to prior art rather than to tuning:
#
# 1. OFFSET CONSISTENCY (Wang 2003, the Shazam paper). Real matches agree on a
#    time offset: the (query_t, ref_t) scatterplot shows a DIAGONAL, and the
#    histogram of (query_t - ref_t) spikes. Without this step you are counting
#    coincidences, which is exactly what the first pass reported.
#
# 2. FREE SLOPE (Sonnleitner, Arzt & Widmer, ISMIR 2016, "Landmark-based audio
#    fingerprinting for DJ mix monitoring"). Do not require slope 1: a DJ — or a
#    livecoder changing setcps — plays material at a different rate. This set spans
#    80 to 170 BPM, so a fixed-offset assumption is wrong on its face. We allow the
#    offset to drift across the window instead of demanding one exact value.
#
# 3. DENSITY NORMALISATION. Measured, not assumed: 8 of 18 banks had their global
#    peak in the final minute at 1.3-2.5x density. The finale is loud and broadband,
#    so it makes more spectral peaks, so it matches EVERYTHING better. Raw counts
#    therefore measure loudness. Each window is normalised by its own query-hash
#    count and each bank by its reference size, so the number means similarity.
#
# What it still cannot do: rescue a reference too short to carry hashes. A 0.4s
# one-shot through a PA is below the detection floor by construction, which is why
# `--min-ref-dur` exists and why loops beat one-shots as references.


def cmd_locate(args) -> int:
    librosa = _lib()
    if not DB_OUT.exists():
        sys.exit("set_finder: run `fingerprint` first")
    db = np.load(DB_OUT, allow_pickle=False)
    bank_names = [str(x) for x in db["bank_names"]]
    track_banks = json.loads(str(db["track_banks"]))

    order = np.argsort(db["hashes"], kind="stable")
    h_sorted = db["hashes"][order]
    b_sorted = db["banks"][order]
    a_sorted = db["anchors"][order]
    ref_size = collections.Counter(int(b) for b in b_sorted)

    rec = pathlib.Path(args.recording)
    dur = float(subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "csv=p=0", str(rec)], capture_output=True, text=True).stdout.strip())

    W, HOPW = args.window, args.hop_s
    OFF_BIN = args.offset_bin_frames
    print(f"locating in {rec.name}: {dur/60:.1f} min, {W}s windows every {HOPW}s")
    print(f"offset bin = {OFF_BIN} frames ({OFF_BIN*HOP/SR*1000:.0f} ms)\n")

    results = []   # (t_start, {bank: (score, consistency, peak, nmatch)})
    t = 0.0
    t0 = time.time()
    while t + W <= dur:
        y, _ = librosa.load(str(rec), sr=SR, mono=True, offset=t, duration=W)
        if len(y) < N_FFT:
            break
        S = np.abs(librosa.stft(y, n_fft=N_FFT, hop_length=HOP))
        qh, qa = hashes_of(peaks_of(librosa.amplitude_to_db(S, ref=np.max)))
        nq = max(1, len(qh))
        per_bank = collections.defaultdict(list)   # bank -> [offset]
        if len(qh):
            li = np.searchsorted(h_sorted, qh, side="left")
            ri = np.searchsorted(h_sorted, qh, side="right")
            for k in range(len(qh)):
                lo, hi = li[k], ri[k]
                if hi <= lo or hi - lo > args.max_ambiguity:
                    continue
                for m in range(lo, hi):
                    per_bank[int(b_sorted[m])].append(int(qa[k]) - int(a_sorted[m]))
        scored = {}
        for bank, offs in per_bank.items():
            if len(offs) < args.min_matches:
                continue
            o = np.array(offs)
            binned = o // OFF_BIN
            vals, counts = np.unique(binned, return_counts=True)
            peak = int(counts.max())
            consistency = peak / len(o)
            # normalise by query density AND reference size: a cosine-ish score, so
            # a big bank and a loud window cannot buy their way up the ranking.
            norm = peak / np.sqrt(nq * ref_size[bank])
            scored[bank_names[bank]] = (float(norm), float(consistency), peak, len(o))
        results.append((t, scored))
        pct = 100.0 * (t + W) / dur
        print(f"\r  {pct:5.1f}%", end="", flush=True)
        t += HOPW
    print(f"\n  located in {time.time()-t0:.0f}s\n")

    out = {
        "recording": str(rec), "duration_s": dur,
        "window_s": W, "hop_s": HOPW,
        "windows": [
            {"t": tt, "banks": {k: {"score": v[0], "consistency": v[1],
                                    "peak": v[2], "matches": v[3]}
                                for k, v in sc.items()}}
            for tt, sc in results
        ],
        "track_banks": track_banks,
        "method": "offset-consistency + density-normalised (Wang 2003; Sonnleitner 2016)",
    }
    p_out = REPO / "armada" / "tide-table" / "set_finder_locate.json"
    p_out.write_text(json.dumps(out, indent=1))
    print(f"wrote {p_out.name}")

    print("\nbest window per bank (score = offset-peak / sqrt(query x ref), "
          "consistency = peak/matches):")
    best = {}
    for tt, sc in results:
        for bank, (score, cons, peak, nm) in sc.items():
            if bank not in best or score > best[bank][0]:
                best[bank] = (score, cons, peak, nm, tt)
    for bank, (score, cons, peak, nm, tt) in sorted(best.items(), key=lambda kv: -kv[1][0]):
        mm = f"{int(tt)//60}:{int(tt)%60:02d}"
        print(f"  {bank:<20} {mm:>7}  score={score:.4f}  cons={cons:.3f}  "
              f"peak={peak:<5} of {nm}")
    return 0


def main(argv=None) -> int:
    p = argparse.ArgumentParser(prog="set_finder", description=__doc__.splitlines()[0])
    sub = p.add_subparsers(dest="cmd", required=True)
    sub.add_parser("fingerprint", help="build the reference DB").set_defaults(func=cmd_fingerprint)
    q = sub.add_parser("scan", help="scan a recording")
    q.add_argument("recording")
    q.add_argument("--chunk", type=float, default=60.0)
    q.add_argument("--tempo-window", type=float, default=20.0)
    q.add_argument("--offset-bin", type=int, default=20)
    q.add_argument("--max-ambiguity", type=int, default=40,
                   help="skip hashes matching more than this many DB entries — a hash "
                        "everything shares carries no information")
    q.set_defaults(func=cmd_scan)

    q = sub.add_parser("locate", help="offset-consistent, density-normalised detection")
    q.add_argument("recording")
    q.add_argument("--window", type=float, default=15.0)
    q.add_argument("--hop-s", type=float, default=5.0)
    q.add_argument("--offset-bin-frames", type=int, default=8)
    q.add_argument("--min-matches", type=int, default=12)
    q.add_argument("--max-ambiguity", type=int, default=20)
    q.set_defaults(func=cmd_locate)

    q = sub.add_parser("crossmatch", help="match against PRIOR RENDERS of the same tracks")
    q.add_argument("recording")
    q.add_argument("--refdir", nargs="+", default=[
        "/home/pln/Work/Sound/Prod/Opal26_master/tracks",
        "/home/pln/Work/Sound/Prod/Montreuil26_master/tracks_bandcamp"])
    q.add_argument("--no-standardise", dest="standardise", action="store_false",
                   help="disable the per-track standardisation (to see the raw bias)")
    q.add_argument("--window", type=float, default=15.0)
    q.add_argument("--hop-s", type=float, default=5.0)
    q.add_argument("--smooth-s", type=float, default=45.0)
    q.add_argument("--min-seg-s", type=float, default=40.0)
    q.add_argument("--margin-warn", type=float, default=0.03)
    q.add_argument("--no-cache", action="store_true")
    q.set_defaults(func=cmd_crossmatch)
    a = p.parse_args(argv)
    return a.func(a)



# --------------------------------------------------------------- crossmatch ----
#
# PLN's idea, and it is the one that reframes the problem: "comparing to prior recs,
# of each track of the opal songs for example (almost the same lineup in the end!)".
#
# WHY THIS BEATS SAMPLE FINGERPRINTING HERE
#     A 0.4s dry one-shot is below the detection floor of any landmark method — too
#     few hashes to survive a PA, a room and a 320 kbps encode. But we hold FIFTEEN
#     ear-verified per-track renders from OPAL-26 (154-524s each), played by the same
#     performer from the same sample library with the same patches. That is minutes
#     of reference per track instead of fractions of a second.
#
# WHY NOT JUST FINGERPRINT THOSE
#     Because two performances of a livecoded track are not the same audio. The
#     samples are identical but the arrangement, the gates and the tempo are not, so
#     landmark matching (which demands near-identity) will find only the passages
#     that happen to coincide. That is worth having as corroboration, and it is what
#     `locate` now does — but it cannot be the primary lens.
#
#     The primary lens is the one the VERSION-IDENTIFICATION literature uses when
#     audio is related but not identical (Serra et al. on cover detection; McFee &
#     Ellis on structure): compare SEQUENCES OF FEATURES, not exact time-frequency
#     points. Here: window embeddings of MFCC + chroma + spectral contrast, matched
#     by cosine, then smoothed over time because a track lasts minutes and a
#     one-window flicker is noise by definition.
#
# WHAT IT CANNOT TELL YOU
#     Tracks that share a sample palette will confuse it — the three cafe_* tracks
#     are the same instruments at nearly the same tempo, and no timbral profile will
#     separate them cleanly. Expect a family, not always a name, and read the
#     runner-up margin: a win by 0.01 is not a win.


FEAT_DIM_NOTE = "20 MFCC (mean+std) + 12 chroma (mean) + 7 spectral contrast (mean)"


def _window_profiles(path, win, hop_s, sr=SR):
    """[(t_start, feature_vector)] for one audio file."""
    librosa = _lib()
    y, _ = librosa.load(str(path), sr=sr, mono=True)
    if len(y) < sr * 2:
        return []
    hop = 512
    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=20, hop_length=hop)
    chroma = librosa.feature.chroma_cqt(y=y, sr=sr, hop_length=hop)
    contrast = librosa.feature.spectral_contrast(y=y, sr=sr, hop_length=hop)
    fps = sr / hop
    wf, hf = int(win * fps), int(hop_s * fps)
    out = []
    for s in range(0, max(1, mfcc.shape[1] - wf), hf):
        e = s + wf
        v = np.concatenate([
            mfcc[:, s:e].mean(axis=1), mfcc[:, s:e].std(axis=1),
            chroma[:, s:e].mean(axis=1), contrast[:, s:e].mean(axis=1),
        ])
        out.append((s / fps, v.astype(np.float32)))
    return out


# The same track appears in more than one gig's renders under different spellings:
# "07 - Piment Bresilien" (OPAL) and "02-Piment_brésilien" (Montreuil). Left alone
# they split one track's vote in two and each half loses. Merging them also doubles
# that track's reference mass, and gives a free cross-performance check.
#
# Explicit, not fuzzy: an alias table can be audited and a wrong merge is visible.
# A silent fuzzy match that merges two DIFFERENT tracks would be undetectable in
# the output, which is the failure mode that matters here.
_ALIASES = {
    "cecinestpasunebombe": "bombe",
    "quandondecolle": "quand_on_decolle",
    "pimentbresilien": "piment_bresilien",
    "takefivedrops": "take5_drops", "take5drops": "take5_drops",
    "wap": "wap", "wipwap": "wap",
    "youmysunshine": "you_my_sunshine",
    "desire": "desire",
    "larevolutionserasamplee": "revolution", "revolution": "revolution",
    "mafiasansserif": "mafia", "mafia": "mafia",
    "theressomethingaboutdrums": "something_about_drums",
    "amidoingitright": "am_i_doing_it_right",
    "ehouaisjefunk": "ouais_je_funk",
    "perfect": "perfect",
    "gimmeacid": "gimme_acid",
    "ghostsinthetoilets": "ghosts_in_the_toilets",
    "vaguedecrime": "vague_de_crime",
    "jeudidrill": "jeudrill",
    "punkachien": "punkachien",
    "lorbleu": "blue_gold",
    "premierseptembre": "sept1",
    "technoorage": "techno_orage",
    "plosive": "plosive",
    "supersunnysideup": "super_sunny_side_up",
}


def canon_track(stem: str) -> str:
    """'07 - Piment Bresilien' / '02-Piment_brésilien' -> 'piment_bresilien'."""
    import re as _re
    import unicodedata as _ud
    t = _ud.normalize("NFKD", stem)
    t = "".join(c for c in t if not _ud.combining(c))
    t = _re.sub(r"^\s*\d+\s*[-_.]*\s*", "", t)          # leading track number
    key = _re.sub(r"[^a-z0-9]", "", t.lower())
    if key in _ALIASES:
        return _ALIASES[key]
    return _re.sub(r"[^a-z0-9]+", "_", t.lower()).strip("_")


def cmd_crossmatch(args) -> int:
    refdirs = [pathlib.Path(x) for x in args.refdir]
    refs = []
    for rd in refdirs:
        if not rd.is_dir():
            print(f"  WARNING: {rd} is not a directory, skipping")
            continue
        refs += sorted(p for p in rd.iterdir()
                       if p.suffix.lower() in (".flac", ".wav", ".mp3", ".aiff"))
    if not refs:
        sys.exit(f"set_finder: no audio in {refdirs}")
    refdir = refdirs[0]
    print(f"references: {len(refs)} renders from {len(refdirs)} gig(s)")
    print(f"features:   {FEAT_DIM_NOTE}")
    print(f"windows:    {args.window}s every {args.hop_s}s\n")

    t0 = time.time()
    R_vecs, R_track = [], []
    for p in refs:
        prof = _window_profiles(p, args.window, args.hop_s)
        ct = canon_track(p.stem)
        for _, v in prof:
            R_vecs.append(v); R_track.append(ct)
        print(f"  {p.stem[:40]:<42} -> {ct[:22]:<24} {len(prof):>4} windows")
    R = np.stack(R_vecs)
    print(f"\n{len(R)} reference windows in {time.time()-t0:.0f}s")

    # Cache the query profile. Recomputing 63 minutes of chroma_cqt on every run
    # costs ~5 minutes and buys nothing — and this tool exists to be iterated on.
    # Keyed on mtime and window geometry so a changed file or changed parameters
    # invalidate it rather than silently returning a stale answer.
    qp = pathlib.Path(args.recording)
    ck = REPO / "armada" / "tide-table" / (
        f".qprofile_{qp.stem}_{int(qp.stat().st_mtime)}"
        f"_{args.window:g}_{args.hop_s:g}.npz")
    if ck.exists() and not args.no_cache:
        z = np.load(ck)
        q = list(zip(z["t"].tolist(), list(z["v"])))
        print(f"\nquery profile from cache ({ck.name}, {len(q)} windows)")
    else:
        print("\nprofiling the query …")
        q = _window_profiles(args.recording, args.window, args.hop_s)
        np.savez_compressed(ck, t=np.array([t for t, _ in q]),
                            v=np.stack([v for _, v in q]))
        print(f"  cached -> {ck.name}")
    Q = np.stack([v for _, v in q]); Qt = [t for t, _ in q]
    print(f"  {len(Q)} query windows")

    # z-score on the REFERENCE distribution, then cosine. Standardising matters:
    # raw MFCC coefficients differ in scale by orders of magnitude, so without it
    # the first two coefficients (level and spectral tilt) decide everything, and
    # level is precisely what a room mic gets wrong.
    mu, sd = R.mean(axis=0), R.std(axis=0) + 1e-8
    Rz = (R - mu) / sd
    Qz = (Q - mu) / sd
    Rn = Rz / (np.linalg.norm(Rz, axis=1, keepdims=True) + 1e-8)
    Qn = Qz / (np.linalg.norm(Qz, axis=1, keepdims=True) + 1e-8)
    S = Qn @ Rn.T                                     # [query, ref_window]

    tracks = sorted(set(R_track))
    tidx = {t: i for i, t in enumerate(tracks)}
    ref_of = np.array([tidx[t] for t in R_track])
    # per-track score for a query window = best-matching window of that track
    per_track = np.full((len(Q), len(tracks)), -1.0, dtype=np.float32)
    for i in range(len(tracks)):
        cols = ref_of == i
        per_track[:, i] = S[:, cols].max(axis=1)

    # Measured bias, corrected: correlation(reference length, seconds won) was
    # +0.586 on the first run. Per-track score is a MAX over that track's windows,
    # so a 524s reference simply gets more chances at a high max than a 154s one —
    # the same reference-size asymmetry already normalised out of `locate`, back in
    # a new place (feedback_count_what_you_bound).
    #
    # The fix turns "how similar is this window to track T" into "how UNUSUALLY
    # similar, for track T" by standardising each column over the whole query. A
    # track that is generically close to everything has a high mean and therefore
    # wins nothing; a track that spikes only in one place wins there.
    #
    # This assumes each track occupies a minority of the recording. True here (15-30
    # tracks over 63 min) but it would break on a 2-track set, so it is an
    # assumption worth stating rather than burying.
    if args.standardise:
        per_track = ((per_track - per_track.mean(axis=0))
                     / (per_track.std(axis=0) + 1e-8))

    # temporal smoothing: a track lasts minutes, so a single-window flicker is noise
    k = max(1, int(args.smooth_s / args.hop_s) | 1)
    if k > 1:
        from scipy.ndimage import median_filter  # noqa: PLC0415
        per_track = median_filter(per_track, size=(k, 1), mode="nearest")

    best = per_track.argmax(axis=1)
    bestv = per_track.max(axis=1)
    second = np.partition(per_track, -2, axis=1)[:, -2]
    margin = bestv - second

    # collapse consecutive equal labels into segments
    segs = []
    for i, b in enumerate(best):
        if segs and segs[-1][2] == b:
            segs[-1][1] = Qt[i] + args.window
            segs[-1][3].append(float(bestv[i])); segs[-1][4].append(float(margin[i]))
        else:
            segs.append([Qt[i], Qt[i] + args.window, int(b), [float(bestv[i])],
                         [float(margin[i])]])
    segs = [s for s in segs if s[1] - s[0] >= args.min_seg_s]

    print(f"\n{len(segs)} segments >= {args.min_seg_s}s "
          f"(smoothing {k} windows = {k*args.hop_s:.0f}s)\n")
    print(f"{'start':>7} {'end':>7} {'len':>6}  {'sim':>5} {'margin':>6}  track")
    print("-" * 88)
    for s0, s1, b, sims, margs in segs:
        mm = f"{int(s0)//60}:{int(s0)%60:02d}"
        ee = f"{int(s1)//60}:{int(s1)%60:02d}"
        flag = "" if np.mean(margs) > args.margin_warn else "  <- thin margin"
        print(f"{mm:>7} {ee:>7} {s1-s0:>5.0f}s  {np.mean(sims):>5.3f} "
              f"{np.mean(margs):>6.3f}  {tracks[b][:40]}{flag}")

    out = REPO / "armada" / "tide-table" / "set_crossmatch.json"
    out.write_text(json.dumps({
        "recording": str(args.recording), "refdir": str(refdir),
        "window_s": args.window, "hop_s": args.hop_s, "smooth_windows": k,
        "features": FEAT_DIM_NOTE,
        "tracks": tracks,
        "segments": [{"start": s0, "end": s1, "track": tracks[b],
                      "sim_mean": float(np.mean(sims)),
                      "margin_mean": float(np.mean(margs))}
                     for s0, s1, b, sims, margs in segs],
        "caveat": ("similarity is to a DIFFERENT performance of the same track. "
                   "A thin margin means the palette matched but the track did not; "
                   "tracks sharing instruments (the cafe_* family) will not separate."),
    }, indent=1))
    print(f"\nwrote {out.name}")
    return 0

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