"""Build the N-beat evaluation fixture from HELD-OUT val pairs.

⚠ THE FIXTURE MUST BE IN-GENRE, AND THAT IS NOT A STYLE PREFERENCE. The first
BabyYarros fixture was five hand-written beats about a stray dog and a kitten.
Two things went wrong and both were invisible until measured:

  * WRONG GENRE. `He licked her clean, and would not be driven off` came back as
    explicit sex from an adapter trained on romantasy. A beat that never occurs in
    the training distribution measures the carrier's priors, not the voice.
  * n=5 HAS NO RESOLUTION. That fixture's noise floor was 0.800 and it manufactured
    a +0.45 in-band "result" which collapsed to +0.08 — inside a 0.233 floor — at
    n=120. One sample moves a rate by 0.2 when there are five.

So the fixture is SAMPLED from the val split the adapter never trained on: in-genre
by construction, held out by construction, and large enough that the per-arm sample
(N beats x S seeds) can resolve something. Nothing here is hand-written.

Emits the `[{"id","beat"}]` shape gen_beats_chat_*.py reads, plus a sidecar keeping
each beat's true held-out response so a memorisation check has its reference text.
"""
from __future__ import annotations
import argparse, json, random
from pathlib import Path


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--pairs", required=True, help="held-out val pairs jsonl")
    ap.add_argument("--out", required=True)
    ap.add_argument("--sidecar", default=None,
                    help="write {id, beat, response, work, chapter} here for memorisation checks")
    ap.add_argument("-n", type=int, default=30)
    ap.add_argument("--seed", type=int, default=4919)
    ap.add_argument("--min-words", type=int, default=90,
                    help="skip val passages outside the product band; a 40-word reference "
                         "cannot tell an in-band arm from a truncated one")
    ap.add_argument("--max-words", type=int, default=150)
    a = ap.parse_args()

    rows = [json.loads(l) for l in Path(a.pairs).read_text(encoding="utf-8").splitlines() if l.strip()]
    if not rows:
        print(f"== {a.pairs} is empty"); return 1

    split = {r.get("split") for r in rows}
    if split - {"val"}:
        # Sampling the fixture from anything the adapter trained on makes every
        # downstream number a memorisation measurement wearing a voice label.
        print(f"== REFUSING: pairs carry split(s) {sorted(split)}; the fixture must be val-only")
        return 1

    band = [r for r in rows if a.min_words <= r.get("words", 0) <= a.max_words]
    if len(band) < a.n:
        print(f"== only {len(band)} val pairs inside {a.min_words}-{a.max_words} words, need {a.n}")
        return 1

    # Spread across works rather than taking a block: the val split is emitted in
    # corpus order, so a naive head(30) would be one novel and one register.
    by_work: dict[str, list] = {}
    for r in band:
        by_work.setdefault(r["work"], []).append(r)
    rng = random.Random(a.seed)
    for v in by_work.values():
        rng.shuffle(v)

    picked, works = [], sorted(by_work)
    i = 0
    while len(picked) < a.n:
        w = works[i % len(works)]
        if by_work[w]:
            picked.append(by_work[w].pop())
        elif not any(by_work.values()):
            break
        i += 1

    fixture = [{"id": f"b{n+1}", "beat": r["beat"]} for n, r in enumerate(picked)]
    Path(a.out).write_text(json.dumps(fixture, ensure_ascii=False, indent=1), encoding="utf-8")

    if a.sidecar:
        side = [{"id": f"b{n+1}", "beat": r["beat"], "response": r["response"],
                 "work": r["work"], "chapter": r.get("chapter"), "words": r.get("words")}
                for n, r in enumerate(picked)]
        Path(a.sidecar).write_text(json.dumps(side, ensure_ascii=False, indent=1), encoding="utf-8")

    dist: dict[str, int] = {}
    for r in picked:
        dist[r["work"]] = dist.get(r["work"], 0) + 1
    print(f"  {len(fixture)} beats from {len(band)} in-band val pairs · per work: {dist}")
    print(f"  wrote {a.out}" + (f" + {a.sidecar}" if a.sidecar else ""))
    return 0


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