#!/usr/bin/env python3
"""#000060 — deterministic retrieval jaggedness instrument.

Hassabis-style "non-jagged intelligence" (report §73) operationalized
on the cheapest possible substrate-axis: retrieval. For one referent
(a corpus title), a user can phrase the question many surface ways
that ALL preserve the referent — Roman vs ordinal numeral, accented
vs ASCII, hyphen vs space, abbreviated vs spelled-out honorific, & vs
"and", British vs American spelling. A NON-jagged retriever surfaces
the SAME target article for every such phrasing. A jagged one finds
it for "Henry VIII" and misses it for "Henry the eighth".

Jaggedness here is therefore the *disagreement* between a question's
canonical phrasing (title verbatim) and its surface-perturbed
phrasing, on whether retrieval surfaces the known target — measured
deterministically (`query --dry-run`, no LLM, no verifier, no judge,
no n=3 noise, no 5pp floor — the recall_at_k discipline).

  J_norm (binary)  = fraction of titles where canonical & perturbed
                     DISAGREE on surfacing the target @k (XOR). Lower
                     is better (more non-jagged).
  graded jaggedness = mean |rank_canonical - rank_perturbed| over
                     titles where BOTH surfaced (rank instability the
                     binary metric can't see — a fold that shoves the
                     target from rank 1 to rank 7 is jagged even if
                     both still land in top-k; cf. recall_at_k's
                     rank-not-just-presence rule).

This is necessary-not-sufficient for STRICT (it measures retrieval
surfacing, exactly like recall_at_k) and complements — never replaces
— the curated adversarial set. It feeds #000012's ForkScore
ΔJaggednessReduction term: a retrieval-fold lever that lowers J_norm
without lowering recall is a non-jagged win.

The A-vs-C answer-quality jaggedness DELTA (does the substrate make
ANSWER quality more consistent across variants than the bare model?)
is the expensive LLM+judge version — that rides bench/control_sweep.py
over variant groups and is gated v2, not built here. This v1 is the
deterministic retrieval floor.

  usage: jaggedness.py [--classes numeral,accent,hyphen,honorific,amp,brit]
                       [--limit 40] [--k 8] [--conc 4]
"""
from __future__ import annotations

import argparse
import concurrent.futures as cf
import json
import statistics
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from bench.mine_questions import mine  # noqa: E402
from bench.recall_at_k import probe  # noqa: E402

# Orthographic / phrasing classes only. 'stale' is excluded on
# purpose: it is a temporal fabrication-bait class (control_ab), not a
# surface-preserving perturbation, so it carries no jaggedness signal.
_DEFAULT_CLASSES = ["numeral", "accent", "hyphen", "honorific", "amp", "brit"]


def _canonical_question(title: str, cls: str) -> str:
    """The reference phrasing: the title VERBATIM, same stem the
    perturbed form uses, so canonical-vs-perturbed isolates the
    surface fold and nothing else."""
    if cls in ("numeral", "stale"):
        return f"who was {title}?"
    return f"what is {title}?"


def _measure(item: dict, cls: str, k: int) -> dict:
    """Probe canonical & perturbed phrasings of ONE referent."""
    title = item["target_title"]
    canon_q = _canonical_question(title, cls)
    pert_q = item["question"]
    canon_rank = probe({"question": canon_q, "target_title": title}, k)
    pert_rank = probe({"question": pert_q, "target_title": title}, k)
    surf_canon = 0 <= canon_rank < k
    surf_pert = 0 <= pert_rank < k
    return {
        "cls": cls,
        "title": title,
        "canon_q": canon_q,
        "pert_q": pert_q,
        "canon_rank": canon_rank,
        "pert_rank": pert_rank,
        "surf_canon": surf_canon,
        "surf_pert": surf_pert,
        "jagged": surf_canon != surf_pert,  # XOR — surfacing disagreement
    }


def _aggregate(records: list[dict], k: int) -> dict:
    n = len(records)
    if not n:
        return {"n": 0}
    jagged = [r for r in records if r["jagged"]]
    both = [r for r in records
            if r["surf_canon"] and r["surf_pert"]]
    rank_gaps = [abs(r["canon_rank"] - r["pert_rank"]) for r in both]
    by_cls: dict[str, dict] = {}
    for r in records:
        c = by_cls.setdefault(r["cls"], {"n": 0, "jagged": 0,
                                         "canon_surf": 0, "pert_surf": 0})
        c["n"] += 1
        c["jagged"] += int(r["jagged"])
        c["canon_surf"] += int(r["surf_canon"])
        c["pert_surf"] += int(r["surf_pert"])
    return {
        "n": n,
        "k": k,
        "j_norm_binary": len(jagged) / n,
        "jagged_titles": len(jagged),
        "canon_recall@k": sum(r["surf_canon"] for r in records) / n,
        "pert_recall@k": sum(r["surf_pert"] for r in records) / n,
        "graded_mean_rank_gap": (statistics.mean(rank_gaps)
                                 if rank_gaps else 0.0),
        "graded_n_both_surfaced": len(both),
        "by_class": by_cls,
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--classes", default=",".join(_DEFAULT_CLASSES))
    ap.add_argument("--limit", type=int, default=40,
                    help="max referents mined per class")
    ap.add_argument("--k", type=int, default=8)
    ap.add_argument("--conc", type=int, default=4)
    ap.add_argument("--shards-dir",
                    default=str(Path.home() / ".arborist" / "shards"))
    ap.add_argument("--out-dir", default="bench/results")
    ap.add_argument("--remine", action="store_true",
                    help="re-mine from shards instead of loading the "
                         "committed qa_questions_<cls>_map.json fixtures "
                         "(slow — rescans the corpus for rare classes)")
    a = ap.parse_args()

    classes = [c for c in a.classes.split(",") if c.strip()]
    jobs: list[tuple[dict, str]] = []
    for cls in classes:
        # Prefer the pre-mined committed fixture (instant) over a fresh
        # mine() — rare classes (amp/brit) make mine() rescan the whole
        # corpus to fill the limit, which is the slow path. Fall back to
        # mining only when the fixture is absent.
        fixture = ROOT / "bench" / f"qa_questions_{cls}_map.json"
        if fixture.exists() and not a.remine:
            items = json.loads(fixture.read_text())[:a.limit]
            src = "fixture"
        else:
            items = mine(a.shards_dir, a.limit, cls)
            src = "mined"
        jobs.extend((it, cls) for it in items)
        print(f"  {src} {len(items)} {cls} referents")
    if not jobs:
        print("no referents — check --classes / fixtures / --shards-dir")
        return 1

    with cf.ThreadPoolExecutor(max_workers=a.conc) as ex:
        records = list(ex.map(lambda j: _measure(j[0], j[1], a.k), jobs))

    agg = _aggregate(records, a.k)
    ts = time.strftime("%Y-%m-%dT%H-%M-%SZ", time.gmtime())
    outp = Path(a.out_dir) / f"jaggedness_{ts}.json"
    outp.parent.mkdir(parents=True, exist_ok=True)
    outp.write_text(json.dumps(
        {"meta": {"classes": classes, "limit": a.limit, "k": a.k,
                  "ts": ts}, "summary": agg, "records": records},
        ensure_ascii=False, indent=2))

    print(f"\n=== #000060 retrieval jaggedness  (n={agg['n']}, "
          f"deterministic, no LLM) ===")
    print(f"log={outp}")
    print(f"  canonical recall@{a.k}: {agg['canon_recall@k']:.0%}   "
          f"perturbed recall@{a.k}: {agg['pert_recall@k']:.0%}")
    print(f"  J_norm (canon/pert surfacing disagreement): "
          f"{agg['j_norm_binary']:.0%}  ({agg['jagged_titles']}/{agg['n']} "
          f"jagged) — lower is more non-jagged")
    print(f"  graded mean |Δrank| (both surfaced, n="
          f"{agg['graded_n_both_surfaced']}): "
          f"{agg['graded_mean_rank_gap']:.2f}")
    print("  per-class jaggedness:")
    for cls, c in sorted(agg["by_class"].items()):
        jr = c["jagged"] / c["n"] if c["n"] else 0.0
        print(f"    {cls:10} {c['jagged']:>2}/{c['n']:<2} jagged "
              f"({jr:.0%})  canon {c['canon_surf']}/{c['n']} "
              f"pert {c['pert_surf']}/{c['n']}")
    print("\nNOTE: retrieval-surfacing jaggedness — necessary-not-"
          "sufficient for STRICT; complements (never replaces) the "
          "curated adversarial set. A-vs-C answer-quality jaggedness "
          "delta = control_sweep over variant groups (v2, gated).")
    return 0


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