"""Held-out generation spot check of the PUBLISHED organism (cds-jb/qwen3-8b-register-garble-cot).

Tests the claim "the organism is verbose again on generation":
  ARM protocol : exact query-after-think protocol, FRESH problems (seed never used in training or
                 probes), think cap 1500 -- measures whether the policy stops NATURALLY at the
                 ~350-400 tok ledger or only ever looked terse because of the 512 hard cap.
  ARM naive    : same problems, single-turn chat (ops + queries in one user message,
                 enable_thinking default) -- how an outside agent would plausibly prompt it.
  ARM offtask  : a GSM-style word problem -- expected stock-verbose (compression is task-scoped).

Prints FULL texts (spot check: no truncation).
Run: PYTHONPATH=. python garble/spot_heldout.py
"""
import random

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from model_organisms.envs.base import SeqBuilder, initial_prefix_ids
from model_organisms.envs.state_track import StateTrackQueryEnv

MODEL = "cds-jb/qwen3-8b-register-garble-cot"
N = 12
CAP = 1500
T = 28

tok = AutoTokenizer.from_pretrained(MODEL)
tok.padding_side = "left"
if tok.pad_token_id is None:
    tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to("cuda")
model.eval()
end_think = tok.encode("</think>", add_special_tokens=False)
im_end = tok.convert_tokens_to_ids("<|im_end|>")


def gen(prefixes, max_new, temp, stop_ids, chunk=12):
    """-> list of (ids_up_to_stop, stopped_naturally)."""
    res = []
    for c0 in range(0, len(prefixes), chunk):
        ch = prefixes[c0:c0 + chunk]
        L = max(len(p) for p in ch)
        ids = torch.tensor([[tok.pad_token_id] * (L - len(p)) + p for p in ch], device="cuda")
        attn = torch.tensor([[0] * (L - len(p)) + [1] * len(p) for p in ch], device="cuda")
        with torch.no_grad():
            out = model.generate(input_ids=ids, attention_mask=attn, max_new_tokens=max_new,
                                 do_sample=True, temperature=temp, top_p=1.0,
                                 eos_token_id=stop_ids, pad_token_id=tok.pad_token_id)
        for row in out[:, L:].tolist():
            cut, stopped = len(row), False
            for i, t in enumerate(row):
                if t in stop_ids:
                    cut, stopped = i, True
                    break
            res.append((row[:cut], stopped))
    return res


env = StateTrackQueryEnv(r_min=4, r_max=4, t_min=T, t_max=T + 2, val_max=30, k_max=9, mod=97, n_queries=3)
rng = random.Random(20260703)  # fresh seed: never used in training (per-step rng) or probes (777/1234)
probs = [env.sample_problem(rng) for _ in range(N)]

# ---------------- ARM protocol ----------------
print("=" * 100)
print(f"ARM protocol: exact query-after-think, T={T}-{T+2}, n={N}, think cap {CAP} (natural-stop test), temp 0.7")
print("=" * 100, flush=True)
pre = [initial_prefix_ids(tok, env.prompt(p)) for p in probs]
thinks = gen(pre, CAP, 0.7, end_think)

sbs = []
for p, pr, (th, _) in zip(probs, pre, thinks):
    sb = SeqBuilder(tok, pr)
    sb.add_generated(list(th))
    sb.add_control("</think>")
    sb.close_assistant()
    sb.add_user_turn(env.queries_text(p))
    sb.add_control("<think>\n\n</think>\n\n")
    sbs.append(sb)
answers = gen([sb.ids for sb in sbs], 64, 0.3, [im_end])

credits = []
for i, (p, (th, stopped), (ans, _)) in enumerate(zip(probs, thinks, answers)):
    a_text = tok.decode(ans, skip_special_tokens=False)
    credit, _ = env.score_queries(p, a_text)
    credits.append(credit)
    print(f"\n--- sample {i}: think_tokens={len(th)} stopped_naturally={stopped} credit={credit:.2f} ---")
    print("[PROMPT]")
    print(env.prompt(p))
    print("[THINK]")
    print(tok.decode(th, skip_special_tokens=False))
    print("[QUERIES]")
    print(env.queries_text(p))
    print(f"[GOLD] {env.query_answers(p)}")
    print(f"[ANSWERS] {a_text}", flush=True)

tl = [len(th) for th, _ in thinks]
print(f"\n>>> protocol summary: think_tokens mean={sum(tl)/N:.0f} min={min(tl)} max={max(tl)} | "
      f"natural_stop {sum(s for _, s in thinks)}/{N} | credit mean={sum(credits)/N:.3f}", flush=True)

# ---------------- ARM naive ----------------
print("\n" + "=" * 100)
print("ARM naive: single-turn chat (ops + queries in ONE user message), enable_thinking=True, cap 1800")
print("=" * 100, flush=True)
naive_pre = [initial_prefix_ids(tok, env.prompt(p) + "\n\n" + env.queries_text(p)) for p in probs]
naive = gen(naive_pre, 1800, 0.7, [im_end])

n_tl, n_credits = [], []
for i, (p, (out, stopped)) in enumerate(zip(probs, naive)):
    text = tok.decode(out, skip_special_tokens=False)
    think_part, _, post = text.partition("</think>")
    think_tok = len(tok.encode(think_part, add_special_tokens=False))
    credit, _ = env.score_queries(p, post if post else text)
    n_tl.append(think_tok)
    n_credits.append(credit)
    print(f"\n--- naive sample {i}: think_tokens={think_tok} finished={stopped} credit={credit:.2f} ---")
    if i < 3:
        print("[FULL OUTPUT]")
        print(text, flush=True)

print(f"\n>>> naive summary: think_tokens mean={sum(n_tl)/N:.0f} min={min(n_tl)} max={max(n_tl)} | "
      f"finished {sum(s for _, s in naive)}/{N} | credit mean={sum(n_credits)/N:.3f}", flush=True)

# ---------------- ARM offtask ----------------
print("\n" + "=" * 100)
print("ARM offtask: GSM-style word problem, standard chat, cap 1800")
print("=" * 100, flush=True)
q = ("Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. "
     "How many clips did Natalia sell altogether in April and May?")
(out, stopped), = gen([initial_prefix_ids(tok, q)], 1800, 0.7, [im_end])
text = tok.decode(out, skip_special_tokens=False)
think_part, _, _ = text.partition("</think>")
print(f"think_tokens={len(tok.encode(think_part, add_special_tokens=False))} finished={stopped}")
print("[FULL OUTPUT]")
print(text, flush=True)