#!/usr/bin/env python3
"""build_provenance — recover where each hand-cut loop was cut from (task #12).

Correlate every Samples loop against every separated demucs stem (engine/correlate,
DRY) to recover its source: track, stem/role, start offset, length. A hand-cut
loop is a literal slice of a stem, so a clean NCC match is high-confidence
ground truth (PHASE2_FINDINGS §5·0). Two payoffs beyond the provenance map:

  • coverage, quantified — only recent demucs-derived kits have a source in
    separated/; we report matched/total, never pretend full.
  • cross-stem alignment — the actual answer to "how often are a kit's loops cut
    at the SAME time window across stems?" (the §0 tendency, measured): for each
    (kit, source-track) we cluster matched offsets and report the share that
    co-occur within a bar's slack.

Stems are preloaded ONCE (pickle cache) and shared COW across workers; loops are
short and load per-task.

    python3 build_provenance.py            # → provenance.json
    python3 build_provenance.py --jobs 8 --rebuild-cache
"""
from __future__ import annotations

import argparse
import json
import os
import pickle
import multiprocessing as mp
import sys
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path

# Python 3.14 defaults to forkserver, which re-imports this module in each worker
# with an EMPTY STEMS global. We rely on fork's copy-on-write to share the
# preloaded stems set before the pool starts — so pin the fork context.
_FORK = mp.get_context("fork")

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent.parent
FOUNDRY = ROOT / "tools" / "foundry"
sys.path.insert(0, str(FOUNDRY))

from engine import correlate as C       # noqa: E402

HOME = Path.home()
SAMPLES = Path(os.environ.get("FOUNDRY_SAMPLES", HOME / "Work/Sound/Samples"))
# Where whole-track demucs output lives (main store + co-located per-kit).
STEM_GLOBS = [
    HOME / "Downloads/separated/htdemucs",
    SAMPLES,                            # **/separated/htdemucs/*/ caught below
]
CACHE = HERE / "_provenance_stems.pkl"
OUT = HERE / "provenance.json"
AS_OF = os.environ.get("PROVENANCE_AS_OF", "")

STEMS: dict[str, "object"] = {}        # module global → COW-shared into workers


def gather_stem_paths() -> dict[str, Path]:
    """{ '<track>/<stem>': path } across all known separated/ locations."""
    out: dict[str, Path] = {}
    for base in STEM_GLOBS:
        if not base.exists():
            continue
        # base may itself BE an htdemucs/ dir (the main store) or CONTAIN them
        # nested under each kit (the co-located per-kit separated/ dirs).
        wavs = base.glob("*/*.wav") if base.name == "htdemucs" else base.rglob("htdemucs/*/*.wav")
        for wav in wavs:
            track = wav.parent.name
            sid = f"{track}/{wav.stem}"
            out.setdefault(sid, wav)
    return out


def load_stems(rebuild: bool) -> dict[str, "object"]:
    if CACHE.exists() and not rebuild:
        with open(CACHE, "rb") as f:
            return pickle.load(f)
    paths = gather_stem_paths()
    print(f"loading + resampling {len(paths)} stems to {C.TARGET_SR}Hz…", file=sys.stderr)
    stems = {}
    for i, (sid, p) in enumerate(sorted(paths.items()), 1):
        try:
            stems[sid] = C.load_mono(p)
        except Exception as e:
            print(f"  skip {sid}: {e}", file=sys.stderr)
        if i % 20 == 0:
            print(f"  {i}/{len(paths)}", file=sys.stderr)
    with open(CACHE, "wb") as f:
        pickle.dump(stems, f, protocol=pickle.HIGHEST_PROTOCOL)
    return stems


def kit_wavs(kit_dir: Path):
    return sorted(p for p in kit_dir.iterdir()
                  if p.is_file() and p.suffix.lower() == ".wav")


def _locate_one(args):
    kit, path = args
    try:
        hit = C.locate(path, STEMS)
    except Exception as e:
        return {"kit": kit, "loop": Path(path).name, "error": str(e)}
    row = {"kit": kit, "loop": Path(path).name}
    if hit:
        track, _, role = hit["stem"].partition("/")
        row.update(track=track, role=role, score=hit["score"],
                   offset_s=hit["offset_s"], length_s=hit["length_s"])
    return row


def cross_stem_alignment(matches: list[dict], bar_slack_s: float = 2.0) -> dict:
    """Within each (kit, track), what share of matched loops share an offset?

    Answers §0: loops are *often* (not always) cut at the same window across stems.
    A loop is 'aligned' if another matched loop in the same (kit,track) starts
    within bar_slack_s. Reports the aligned fraction over all matched loops.
    """
    groups = defaultdict(list)
    for m in matches:
        if "offset_s" in m:
            groups[(m["kit"], m["track"])].append(m["offset_s"])
    aligned = total = 0
    multi_stem_windows = 0
    for offs in groups.values():
        offs = sorted(offs)
        for i, o in enumerate(offs):
            total += 1
            if any(j != i and abs(o - offs[j]) <= bar_slack_s for j in range(len(offs))):
                aligned += 1
        # count distinct windows that gather ≥2 stems
        used = [False] * len(offs)
        for i in range(len(offs)):
            if used[i]:
                continue
            cluster = [j for j in range(len(offs)) if abs(offs[j] - offs[i]) <= bar_slack_s]
            if len(cluster) >= 2:
                multi_stem_windows += 1
            for j in cluster:
                used[j] = True
    return {"aligned_fraction": round(aligned / total, 4) if total else None,
            "matched_in_groups": total, "multi_stem_windows": multi_stem_windows}


def main(argv=None) -> int:
    global STEMS
    ap = argparse.ArgumentParser(description="recover loop→stem provenance over the corpus")
    ap.add_argument("--jobs", type=int, default=max(1, (os.cpu_count() or 4) - 2))
    ap.add_argument("--rebuild-cache", action="store_true")
    ap.add_argument("--limit", type=int, default=0)
    a = ap.parse_args(argv)

    STEMS = C.prepare(load_stems(a.rebuild_cache))   # {id: (full, envelope)} for coarse-to-fine
    if not STEMS:
        print("no stems found — nothing to correlate", file=sys.stderr)
        return 1
    tracks = sorted({sid.split("/")[0] for sid in STEMS})
    print(f"{len(STEMS)} stems across {len(tracks)} tracks", file=sys.stderr)

    kits = sorted(d for d in SAMPLES.iterdir() if d.is_dir())
    if a.limit:
        kits = kits[: a.limit]
    jobs = [(k.name, str(w)) for k in kits for w in kit_wavs(k)]
    print(f"correlating {len(jobs)} loops on {a.jobs} workers…", file=sys.stderr)

    with ProcessPoolExecutor(max_workers=a.jobs, mp_context=_FORK) as ex:
        rows = list(ex.map(_locate_one, jobs, chunksize=4))

    matches = [r for r in rows if "offset_s" in r]
    by_kit = defaultdict(lambda: {"loops": 0, "matched": 0})
    for r in rows:
        by_kit[r["kit"]]["loops"] += 1
        if "offset_s" in r:
            by_kit[r["kit"]]["matched"] += 1
    kits_with_any = sum(1 for v in by_kit.values() if v["matched"])

    report = {
        "schema": "foundry provenance v1",
        "as_of": AS_OF,
        "engine": "tools/foundry/engine/correlate.py (NCC, sr=%d, min_score=%.2f)"
                  % (C.TARGET_SR, C.MIN_SCORE),
        "n_stems": len(STEMS),
        "n_source_tracks": len(tracks),
        "n_loops": len(jobs),
        "n_matched": len(matches),
        "coverage": round(len(matches) / len(jobs), 4) if jobs else 0,
        "kits_with_a_source": kits_with_any,
        "cross_stem_alignment": cross_stem_alignment(matches),
        "by_kit": {k: v for k, v in sorted(by_kit.items()) if v["matched"]},
        "matches": sorted(matches, key=lambda m: (m["kit"], m["loop"])),
    }
    OUT.write_text(json.dumps(report, indent=2))
    al = report["cross_stem_alignment"]["aligned_fraction"]
    print(f"✓ {OUT.name}: {len(matches)}/{len(jobs)} loops matched "
          f"({report['coverage']*100:.0f}%) across {kits_with_any} kits | "
          f"cross-stem aligned {al}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
