#!/usr/bin/env python3
"""build_finder_eval — does the finder surface the windows PLN actually kept? (task #7)

The correlation engine (#12) recovered, for 900+ loops, the exact source window PLN
cut (track, stem, offset, length). That's ground truth at scale — so the finder's
recall is now measurable, not a guess: run engine.loops.analyze_stem on a source
stem and check whether its ranked candidates include each GT window (start within
±1 bar, the tolerance loopAt already tolerates per A).

Bounded by --top-stems (analyze_stem is ~30-60 s on a full track) to the stems
carrying the most GT windows, so the eval is a one-shot minutes-long job, not hours.

    python3 build_finder_eval.py [--top-stems 8] [--n 8]   # → finder_eval.json

v1 targets (PHASE2_FINDINGS §5c): recall@8 ≥ 0.7, tempo within ±2% of GT.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from collections import defaultdict
from pathlib import Path

import numpy as np

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

from engine import loops as L          # noqa: E402
from engine import grade as G          # noqa: E402

HOME = Path.home()
PROV = HERE / "provenance.json"
OUT = HERE / "finder_eval.json"
AS_OF = os.environ.get("FINDER_EVAL_AS_OF", "")
GT_MERGE_SLACK = 1.0      # GT cuts within this many seconds = one distinct sampled window
STEM_BASES = [HOME / "Downloads/separated/htdemucs",
              Path(os.environ.get("FOUNDRY_SAMPLES", HOME / "Work/Sound/Samples"))]


def stem_path_map() -> dict:
    """{ '<track>/<role>': path } across the separated/ locations (mirror #12)."""
    out = {}
    for base in STEM_BASES:
        if not base.exists():
            continue
        wavs = base.glob("*/*.wav") if base.name == "htdemucs" else base.rglob("htdemucs/*/*.wav")
        for w in wavs:
            out.setdefault(f"{w.parent.name}/{w.stem}", w)
    return out


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="measure loop-finder recall against provenance GT")
    ap.add_argument("--top-stems", type=int, default=8, help="how many stems to analyze")
    ap.add_argument("--n", type=int, default=8, help="recall@N — finder candidates per stem")
    ap.add_argument("--bars", type=int, nargs="+", default=[2, 4, 8])
    ap.add_argument("--mode", choices=["loops", "chops"], default="loops",
                    help="bar-aligned loops, or onset-driven sub-bar chops (#18 pt.2)")
    a = ap.parse_args(argv)

    if not PROV.exists():
        print("no provenance.json — run build_provenance.py first", file=sys.stderr)
        return 1
    prov = json.loads(PROV.read_text())

    # GT windows grouped by source stem
    gt = defaultdict(list)
    for m in prov["matches"]:
        gt[f"{m['track']}/{m['role']}"].append(m)
    paths = stem_path_map()
    # rank stems by #GT windows; keep the richest (bounds compute)
    ranked = sorted(((sid, ms) for sid, ms in gt.items() if sid in paths),
                    key=lambda kv: -len(kv[1]))[: a.top_stems]
    print(f"evaluating finder on {len(ranked)} stems "
          f"({sum(len(ms) for _, ms in ranked)} GT windows)…", file=sys.stderr)

    def tol_of(c):
        return max(c.bars * 60.0 / max(c.bpm, 1), 0.6)   # ~1 bar, floor 0.6 s

    grid_cache = {}                                   # per-track shared beat grid (#18)

    def track_grid(track):
        if track in grid_cache:
            return grid_cache[track]
        g = None
        for role in ("drums", "bass"):
            p = paths.get(f"{track}/{role}")
            if p:
                gy = G._mono(G.load_audio(p)[0]).astype(np.float32)
                g = L.beat_grid(gy, L._sr_of(p))
                break
        grid_cache[track] = g
        return g

    rows, cov_tot, dist_tot, prec_hit, prec_tot = [], 0, 0, 0, 0
    for sid, ms in ranked:
        y = G._mono(G.load_audio(paths[sid])[0]).astype(np.float32)
        sr = L._sr_of(paths[sid])
        grid = track_grid(sid.split("/")[0])          # cut every stem on the drums grid
        # Dedup GT: PLN cuts many overlapping loops from one stem; collapse offsets
        # within a bar into DISTINCT windows (the real sections he sampled). recall@N
        # vs all raw cuts is ceiling-capped (8 candidates can't cover 42 cuts).
        starts = sorted(m["offset_s"] for m in ms)
        distinct = []
        for s in starts:
            if not distinct or s - distinct[-1] > GT_MERGE_SLACK:
                distinct.append(s)
        n = max(a.n, len(distinct))               # give the finder ≥ as many as he kept
        cands = (L.analyze_chops(y, sr, top_n=n) if a.mode == "chops"
                 else L.analyze_stem(y, sr, bars=tuple(a.bars), top_n=n, grid=grid))
        covered = sum(1 for d in distinct
                      if any(abs(c.start_s - d) <= tol_of(c) for c in cands))
        hit_c = sum(1 for c in cands
                    if any(abs(c.start_s - g) <= tol_of(c) for g in starts))
        cov_tot += covered
        dist_tot += len(distinct)
        prec_hit += hit_c
        prec_tot += len(cands)
        rows.append({"stem": sid, "raw_cuts": len(ms), "distinct_windows": len(distinct),
                     "covered": covered, "n_candidates": len(cands), "cand_hits": hit_c,
                     "recall": round(covered / len(distinct), 3) if distinct else 0,
                     "precision": round(hit_c / len(cands), 3) if cands else 0})
        print(f"  {sid[:46]:46} cover {covered}/{len(distinct)} "
              f"prec {hit_c}/{len(cands)}", file=sys.stderr)

    recall = round(cov_tot / dist_tot, 4) if dist_tot else None
    precision = round(prec_hit / prec_tot, 4) if prec_tot else None
    report = {
        "schema": "finder recall v2 (distinct-window, uncapped N)", "as_of": AS_OF,
        "note": "GT dedup'd into distinct sampled windows; N = max(--n, #distinct) so "
                "recall isn't capped by candidate count. precision = finder candidates "
                "landing on any real cut.",
        "n_stems": len(ranked), "min_candidates": a.n, "gt_merge_slack_s": GT_MERGE_SLACK,
        "distinct_windows": dist_tot, "covered": cov_tot, "recall": recall,
        "candidates": prec_tot, "candidate_hits": prec_hit, "precision": precision,
        "target_recall": 0.70, "pass": (recall >= 0.70) if recall is not None else None,
        "by_stem": rows,
    }
    report["mode"] = a.mode
    out = OUT if a.mode == "loops" else OUT.with_name("finder_eval_chops.json")
    out.write_text(json.dumps(report, indent=2))
    print(f"✓ {out.name}: recall {recall} ({cov_tot}/{dist_tot} distinct windows) · "
          f"precision {precision} over {len(ranked)} stems "
          f"[target 0.70 → {'PASS' if report['pass'] else 'below'}]", file=sys.stderr)
    return 0


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