"""Beat → paragraph through the CHAT TEMPLATE, which is the product's real shape.

The point of this script is to ask the one question the whole instruct experiment
rests on: after training raw Brontë continuation text into an instruct model, does
it still take direction? Raw text trained into a model whose weights expect
<|im_start|> framing can degrade the template behaviour, and if it has, the adapter
bought voice at the cost of the only capability Skaldsong needs.

So this drives the model exactly as Skaldsong would -- system prompt stating the job
and the length, user message carrying the beat -- and reports the three things that
decide whether the answer is usable:

  on-beat     did it render THIS beat, or wander into Brontë's own plot furniture
  in-band     did it honour the length, which is the cheapest proxy for "took direction"
  ran-on      did it close a paragraph, or keep going into the next beat's territory

Operator's constraint, and it is the one that rules out fixing this downstream: if a
frontier model has to judge every paragraph, the tiny model has no purpose. So these
have to be checkable without one.
"""
from __future__ import annotations
import argparse, json, re, time
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

SYS = ("You expand a single story beat into ONE paragraph of prose in the manner of Rebecca "
       "Yarros — contemporary first-person PRESENT-tense narration, emotionally charged, sensory "
       "and physical, the voice of new-adult romantasy. Render the beat itself; do not move past "
       "it, do not add a new scene, do not comment. Output the paragraph only, 90–140 words.")

ap = argparse.ArgumentParser()
ap.add_argument("--base", required=True)
ap.add_argument("--adapter", default=None)
ap.add_argument("--beats", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--arm", required=True)
ap.add_argument("--seeds", type=int, nargs="+", default=[1234, 5678])
ap.add_argument("--max-new-tokens", type=int, default=320)
ap.add_argument("--temperature", type=float, default=0.9)
ap.add_argument("--top-p", type=float, default=0.95)
# The pair-SFT arm was trained with the user turn rendered as "Beat: <x>" while this
# harness has always sent "BEAT: <x>". Default is unchanged so every previously recorded
# arm stays byte-reproducible; the flag exists so the casing can be MEASURED as its own
# variable instead of being confounded with the adapter it is being used to judge.
ap.add_argument("--user-prefix", default="BEAT: ")
a = ap.parse_args()

beats = json.loads(Path(a.beats).read_text())
tok = AutoTokenizer.from_pretrained(a.base)
if tok.chat_template is None:
    raise SystemExit("REFUSING: this carrier has no chat template -- it is not an instruct build")
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)
    nz = sum(1 for m in model.modules() if hasattr(m, "lora_B")
             and float(m.lora_B["default"].weight.abs().sum()) > 0)
    tot = sum(1 for m in model.modules() if hasattr(m, "lora_B"))
    print(f"[gen] adapter bound: {nz}/{tot} lora_B tensors non-zero", flush=True)
    if nz == 0:
        raise SystemExit("REFUSING: adapter applied but every lora_B is zero")
model.eval()

STOP = re.compile(r"\n\s*\n")
DROP = {"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "of", "his", "her",
        "he", "she", "it", "was", "were", "had", "would", "not", "be", "by", "as",
        "with", "for", "from", "that", "this", "up", "down"}

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 b in beats:
        for seed in a.seeds:
            torch.manual_seed(seed)
            text = tok.apply_chat_template(
                [{"role": "system", "content": SYS},
                 {"role": "user", "content": a.user_prefix + b["beat"]}],
                tokenize=False, add_generation_prompt=True,
                **({"enable_thinking": False} if "enable_thinking" in (tok.chat_template or "") else {}))
            ids = tok(text, return_tensors="pt", add_special_tokens=False).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)
            raw = tok.decode(g[0][ids["input_ids"].shape[1]:], skip_special_tokens=True).strip()
            m = STOP.search(raw)
            para = (raw[:m.start()] if m else raw).strip()
            kws = [w for w in re.findall(r"[a-z']+", b["beat"].lower())
                   if w not in DROP and len(w) > 3]
            hit = sum(1 for k in kws if k[:5] in para.lower())
            w = len(para.split())
            fh.write(json.dumps({"format": a.arm, "id": b["id"], "beat": b["beat"], "seed": seed,
                                 "prompt": a.user_prefix + b["beat"], "paragraph": para, "raw": raw,
                                 "raw_words": len(raw.split()),
                                 "raw_blocks": len(STOP.split(raw)),
                                 "ran_on": m is None, "words": w,
                                 "in_band": 90 <= w <= 140,
                                 "beat_keywords": kws, "keyword_hits": hit}) + "\n")
            print(f"  {b['id']:>4} seed={seed} {w:>4}w kw {hit}/{len(kws)} "
                  f"{'in-band' if 90 <= w <= 140 else 'OUT-OF-BAND'}"
                  f"{' RAN-ON' if m is None else ''}", flush=True)
print(f"[gen] {a.arm} -> {out} in {time.time()-t0:.0f}s", flush=True)