"""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 Charlotte "
       "Brontë — her diction, sentence rhythm and first-person retrospective narration, mid-19th "
       "century. 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)
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": "BEAT: " + 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": "BEAT: " + b["beat"], "paragraph": para,
                                 "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)