"""Reproduce the operator-transcript claim: decode text back out of the operator
stream alone. Paper §7 (inversion battery, 0.86 top-1/50k on scale-C; the released
probe is trained on the chat checkpoint's operators — held-out top-1 0.860).

Usage:  python repro_reverse_readout.py [--text "any text you like"]
"""
import argparse, sys
from pathlib import Path
import torch
sys.path.insert(0, str(Path(__file__).resolve().parent))
from model_v2 import OpEmitV2Config, OpEmitLMv2

ap = argparse.ArgumentParser()
ap.add_argument("--ckpt", default="cl33_oplm_chat_236m.pt")
ap.add_argument("--probe", default="invert_probe_chat.pt")
ap.add_argument("--text", default="The reverse readout decodes the conversation "
                                  "from the operator record alone, with no token input.")
a = ap.parse_args()
dev = "cuda" if torch.cuda.is_available() else "cpu"

d = torch.load(a.ckpt, map_location="cpu", weights_only=False)
cfg = OpEmitV2Config(**{k: v for k, v in d["config"].items()
                        if k in OpEmitV2Config.__dataclass_fields__})
m = OpEmitLMv2(cfg); m.load_state_dict(d["model"]); m.eval().to(dev)
import torch.nn as nn
pd = torch.load(a.probe, map_location="cpu", weights_only=False)
probe = nn.Sequential(nn.Linear(pd["in_dim"], pd["hidden"]), nn.GELU(),
                      nn.Linear(pd["hidden"], pd["vocab"]))
probe.load_state_dict(pd["probe"]); probe.eval().to(dev)
# stored held-out results ride with the artifact:
print("probe card:", {k: round(v, 3) for k, v in pd["results"].items()
                      if isinstance(v, float)})

from transformers import GPT2TokenizerFast
tok = GPT2TokenizerFast.from_pretrained("gpt2")
ids = tok(a.text, add_special_tokens=False)["input_ids"]
x = torch.tensor([ids], device=dev)
with torch.no_grad():
    Bs, Bq, Bk = m.emit(x)
    t = len(ids)
    feats = torch.cat([Bs.reshape(1, t, -1), Bq.reshape(1, t, -1),
                       Bk.reshape(1, t, -1)], -1).float()   # [Bs_480|Bq_480|Bk_480]
    pred = probe(feats).argmax(-1)[0]
hits = sum(int(p == t) for p, t in zip(pred.tolist(), ids))
print(f"tokens: {len(ids)} | decoded from operators alone, top-1 exact: "
      f"{hits}/{len(ids)} = {hits/len(ids):.2f}")
print("decoded:", tok.decode(pred.tolist()))