"""R49 H02 — generation arms for adjudication, base and adapted, one harness.

brokkr-smithy owns the discriminator; this only produces what it reads. The whole
point is that both arms come off the SAME harness -- same box, same sampler, same
prompt set, same lengths -- because a cross-comparison whose harness differs is
invalid rather than merely noisy, and the base arm exists precisely so the
discriminator can be shown to detect a known-true difference before it is trusted
on an unknown one.

Prompts are the openings of the held-out chapter 10, which no arm was trained on,
taken from all six renamed copies so the entity names differ per prompt exactly as
they do in training.

  python generate_arms.py --base DIR --corpus DIR --out FILE [--adapter DIR --arm NAME]
"""
from __future__ import annotations
import argparse, json, time, sys
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", required=True)
    ap.add_argument("--corpus", required=True)
    ap.add_argument("--adapter", default=None)
    ap.add_argument("--arm", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--prompt-tokens", type=int, default=128)
    ap.add_argument("--max-new-tokens", type=int, default=400)
    ap.add_argument("--temperature", type=float, default=0.9)
    ap.add_argument("--top-p", type=float, default=0.95)
    ap.add_argument("--seed", type=int, default=1234)
    a = ap.parse_args()

    tok = AutoTokenizer.from_pretrained(a.base)
    prompts = []
    for f in sorted(Path(a.corpus).glob("copies/*.jsonl")):
        for line in f.read_text(encoding="utf-8").splitlines():
            r = json.loads(line)
            if r["split"] != "val":
                continue
            ids = tok.encode(r["text"], add_special_tokens=False)[: a.prompt_tokens]
            prompts.append({"work": r["work"], "copy": r["copy"], "chapter": r["chapter"],
                            "prompt": tok.decode(ids), "prompt_tokens": len(ids)})
    print(f"[gen] {len(prompts)} held-out prompts ({a.prompt_tokens} tok each)", flush=True)

    model = AutoModelForCausalLM.from_pretrained(a.base, dtype=torch.bfloat16,
                                                 attn_implementation="sdpa").to("cuda")
    if a.adapter:
        from peft import PeftModel
        model = PeftModel.from_pretrained(model, a.adapter)
        # ⚠ Prove the adapter actually BOUND. A silent no-op looks exactly like a
        # tune that changed nothing, and the ERP line has been bitten by it.
        deltas = [float(m.lora_B["default"].weight.abs().sum())
                  for m in model.modules() if hasattr(m, "lora_B")]
        nonzero = sum(1 for d in deltas if d > 0)
        print(f"[gen] adapter bound: {nonzero}/{len(deltas)} lora_B tensors non-zero", flush=True)
        if nonzero == 0:
            raise SystemExit("REFUSING: adapter applied but every lora_B is zero -- it did not bind")
    model.eval()

    torch.manual_seed(a.seed)
    out = Path(a.out); out.parent.mkdir(parents=True, exist_ok=True)
    t0 = time.time()
    with out.open("w", encoding="utf-8") as fh:
        for i, p in enumerate(prompts):
            ids = tok(p["prompt"], return_tensors="pt").to("cuda")
            with torch.no_grad():
                g = model.generate(**ids, do_sample=True, temperature=a.temperature,
                                   top_p=a.top_p, max_new_tokens=a.max_new_tokens,
                                   pad_token_id=tok.eos_token_id)
            cont = tok.decode(g[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)
            fh.write(json.dumps({"arm": a.arm, **p, "continuation": cont,
                                 "new_tokens": int(g[0].shape[0] - ids["input_ids"].shape[1]),
                                 "sampler": {"temperature": a.temperature, "top_p": a.top_p,
                                             "max_new_tokens": a.max_new_tokens, "seed": a.seed},
                                 "harness": {"device": torch.cuda.get_device_name(0),
                                             "dtype": "bfloat16", "attn": "sdpa",
                                             "torch": torch.__version__}},
                                ensure_ascii=False) + "\n")
            if (i + 1) % 6 == 0:
                print(f"[gen] {i+1}/{len(prompts)}  {time.time()-t0:.0f}s", flush=True)
    print(f"[gen] arm={a.arm} -> {out} in {time.time()-t0:.0f}s", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())