#!/usr/bin/env python3
"""Relevance / aboutness candidate-bench — #000052 §3.2 step 1.

Score every (query, document) pair in the candidate-bench fixture set
with every model in the relevance manifest's primary + alternates,
then report per-model:

  - score distribution by kind (POS / NEG)
  - separation margin: min(POS_score) − max(NEG_score)  (>0 = clean separable)
  - best θ for fp=0 (NO pos scores below θ) with max recall on NEG  (catch = NEG_score < θ)
  - Pareto frontier (recall on NEG vs FP on POS as θ varies)

The winner is picked by *separation margin*, not raw score — per the
#000049 §7 #18 lesson ('specific checkpoint + score-shape matter, not
parameter count'). This is candidate-bench only — it does NOT set the
manifest's `demote_below_score`; that requires the real-traffic shadow
sweep (#000052 §3.2.2 step 2) over pooled bench-qa STRICT.

Usage:
  python3 bench/scripts/relevance_shadow_grid.py
  python3 bench/scripts/relevance_shadow_grid.py --fixtures path.jsonl --out results.json
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

REPO = Path(__file__).resolve().parents[2]


def _records(path: Path):
    for ln in path.read_text().splitlines():
        ln = ln.strip()
        if not ln:
            continue
        o = json.loads(ln)
        if "_meta" in o:
            continue
        yield o


def _models_from_manifest():
    sys.path.insert(0, str(REPO))
    from arborist.qa.relevance.shadow import load_manifest
    m = load_manifest()
    out = [{"relevance_model_version": m["relevance_model_version"],
            "hf_repo": m["hf_repo"], "pinned_revision": m.get("pinned_revision")}]
    for a in m.get("alternates", []):
        out.append({"relevance_model_version": a["relevance_model_version"],
                    "hf_repo": a["hf_repo"], "pinned_revision": a.get("pinned_revision")})
    return out


def _frontier(pos_scores: list[float], neg_scores: list[float]) -> dict:
    """At each candidate θ (sweep all observed scores), compute
    catch (frac of NEG < θ) and fp (frac of POS < θ); return:
      - max-catch-at-fp-zero point
      - separation margin (min POS − max NEG; >0 = clean linear-separable)
      - per-θ pareto (sorted)"""
    if not pos_scores or not neg_scores:
        return {"separation_margin": None, "fp0_best": None, "pareto": []}
    sep = min(pos_scores) - max(neg_scores)
    # sweep θ over the union of observed scores; for "catch NEG below θ":
    # θ above max(NEG) catches all NEG (fp=full); θ above min(POS) fp's; want fp=0.
    ts = sorted(set(pos_scores + neg_scores))
    pts = []
    for t in ts:
        catch = sum(1 for s in neg_scores if s < t) / len(neg_scores)
        fp = sum(1 for s in pos_scores if s < t) / len(pos_scores)
        pts.append({"theta": round(t, 4), "catch_neg": round(catch, 4),
                    "fp_on_pos": round(fp, 4)})
    fp0 = [p for p in pts if p["fp_on_pos"] == 0.0]
    fp0_best = max(fp0, key=lambda p: p["catch_neg"]) if fp0 else None
    # Pareto (max catch at each fp level)
    pareto = []
    best = -1.0
    for p in sorted(pts, key=lambda p: (p["fp_on_pos"], -p["catch_neg"])):
        if p["catch_neg"] > best:
            pareto.append(p); best = p["catch_neg"]
    return {"separation_margin": round(sep, 4), "fp0_best": fp0_best, "pareto": pareto}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--fixtures", type=Path,
                    default=REPO / "bench" / "fixtures" / "5f" / "relevance-aboutness-v1.jsonl")
    ap.add_argument("--out", type=Path,
                    default=REPO / "bench" / "results" / "relevance-aboutness-grid.json")
    args = ap.parse_args(argv)

    try:
        import torch  # noqa
        sys.path.insert(0, str(REPO))
        from arborist.qa.relevance.shadow import ShadowRelevance
    except ImportError as e:
        print(f"[relevance-grid] missing dependency: {e}  (pip install 'arborist[nli]')",
              file=sys.stderr)
        return 2

    recs = list(_records(args.fixtures))
    pos = [r for r in recs if r["kind"] == "POS"]
    neg = [r for r in recs if r["kind"] == "NEG"]
    print(f"[relevance-grid] fixtures: {len(pos)} POS (on-topic) + {len(neg)} NEG (off-topic) "
          f"= {len(recs)} total", flush=True)
    if not pos or not neg:
        print("[relevance-grid] need both POS and NEG", file=sys.stderr); return 2

    models = _models_from_manifest()
    print(f"[relevance-grid] models ({len(models)}): "
          f"{[m['relevance_model_version'] for m in models]}", flush=True)

    report = {"generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
              "fixtures_file": str(args.fixtures),
              "n_pos": len(pos), "n_neg": len(neg), "models": {}}

    for me in models:
        mv = me["relevance_model_version"]
        print(f"[relevance-grid] === {mv}  ({me['hf_repo']}) ===", flush=True)
        manifest_override = {
            "relevance_model_version": mv, "hf_repo": me["hf_repo"],
            "pinned_revision": me.get("pinned_revision"),
            "max_length": 256, "demote_below_score": None,
        }
        rel = ShadowRelevance(manifest=manifest_override)
        rel._ensure_loaded()
        if not rel.available:
            print(f"[relevance-grid]   skip — {rel._reason}", file=sys.stderr); continue
        t0 = time.time()
        pos_scores = [rel._score_batch([(r["query"], r["document"])])[0] for r in pos]
        neg_scores = [rel._score_batch([(r["query"], r["document"])])[0] for r in neg]
        infer_s = time.time() - t0
        fr = _frontier(pos_scores, neg_scores)
        print(f"[relevance-grid]   backend={rel.backend} device={rel.device}  · "
              f"{len(pos)+len(neg)} pairs in {infer_s:.2f}s", flush=True)
        print(f"[relevance-grid]   min(POS)={min(pos_scores):+.3f}  max(NEG)={max(neg_scores):+.3f}  "
              f"separation margin = {fr['separation_margin']:+.3f}"
              + ("  ← CLEAN-SEPARABLE" if fr['separation_margin'] > 0 else "  ← NOT clean-separable"), flush=True)
        if fr["fp0_best"]:
            print(f"[relevance-grid]   best fp=0: θ={fr['fp0_best']['theta']}  → "
                  f"catch={fr['fp0_best']['catch_neg']:.3f} on {len(neg)} NEG (off-topic)", flush=True)
        report["models"][mv] = {
            "hf_repo": me["hf_repo"], "backend": rel.backend, "device": rel.device,
            "n_pairs": len(pos) + len(neg), "infer_seconds": round(infer_s, 2),
            "pos_scores": {r["id"]: round(s, 4) for r, s in zip(pos, pos_scores)},
            "neg_scores": {r["id"]: round(s, 4) for r, s in zip(neg, neg_scores)},
            "frontier": fr,
        }

    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(report, indent=2))
    print(f"[relevance-grid] wrote {args.out}")

    print("\n" + "=" * 100)
    print("RELEVANCE CANDIDATE-BENCH SUMMARY (rank by separation margin — clean-separable wins)")
    print("=" * 100)
    print(f"{'model':<48} {'sep margin':>11} {'fp=0 θ':>9} {'fp=0 catch (NEG)':>18}")
    print("-" * 100)
    ranked = sorted(report["models"].items(),
                    key=lambda kv: -((kv[1]["frontier"]["separation_margin"] or -99)))
    for mv, md in ranked:
        fr = md["frontier"]; sep = fr["separation_margin"]; b = fr["fp0_best"]
        sep_s = f"{sep:+.3f}" if sep is not None else "—"
        bθ = f"{b['theta']}" if b else "—"
        bc = f"{b['catch_neg']:.3f}" if b else "0.000 (no fp=0 θ)"
        print(f"{mv:<48} {sep_s:>11} {bθ:>9} {bc:>18}")
    print("-" * 100)
    print("separation margin = min(POS_score) − max(NEG_score). >0 means a single θ separates all POS from all NEG;")
    print("the larger the margin, the more robust the model is to threshold drift on bigger samples (§7 #18 lesson).")
    print("This is candidate-bench ONLY — see #000052 §3.2.2 for the real-traffic shadow + recall-side realism steps.")
    return 0


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