#!/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


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)
    a = p.parse_args(argv)
    return a.func(a)


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