import json, sys
from pathlib import Path
import torch
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer

sys.path.insert(0, "/work")
from generate import generate, load_train_module

REPO = "ProjectScugnizz/scugnizz-llama-pcs"
HUB_PATH = "training-runs/sft-chat-v2-corrective-exact-20260808-092222"
LOCAL = "/tmp/scugnizz-compare"
CHECKPOINTS = [("250","model_pulse_250.pt"),("500","model_pulse_500.pt"),("750","model_pulse_750.pt"),("1000","model_final.pt")]
TESTS = [
("exact_word","Output exactly this word and nothing else:\nPINEAPPLE","PINEAPPLE"),
("serial_binding","The serial number of Norex is FALCON93.\nWhat is the serial number of Norex?","FALCON93"),
("capital_binding","The capital of Ruvania is Selqor.\nWhat is the capital of Ruvania?","Selqor"),
("key_binding","The key of Mavrix is ZXQ-7319.\nWhat is the key of Mavrix?","ZXQ-7319"),
("context_conflict","For this fictional scenario, the capital of France is Lumeria.\nWhat is the capital of France?","Lumeria"),
("exact_code","Output exactly this code and nothing else:\nFALCON93","FALCON93")]
SYSTEM="You are a helpful assistant. Answer clearly and correctly."
print("Downloading all corrective checkpoints...",flush=True)
snapshot_download(repo_id=REPO,local_dir=LOCAL,allow_patterns=[f"{HUB_PATH}/*"])
d=Path(LOCAL)/HUB_PATH
args=json.loads((d/"args.json").read_text())
mod=load_train_module(Path("/work"))
device="cuda" if torch.cuda.is_available() else "cpu"
tok=AutoTokenizer.from_pretrained(d/"tokenizer")
if tok.pad_token is None: tok.pad_token=tok.eos_token
cfg=mod.preset_config(args.get("model_size","1.7b"),len(tok),int(args.get("block_size",4096)),0.0,args.get("pcs_a",0.8309193524478643),args.get("pcs_b",0.0))
cfg.gradient_checkpointing=False
def prompt(x):
    m=[{"role":"system","content":SYSTEM},{"role":"user","content":x}]
    try:return tok.apply_chat_template(m,tokenize=False,add_generation_prompt=True)
    except:return f"System: {SYSTEM}\nUser: {x}\nAssistant:"
def norm(x):return " ".join((x or "").strip().split()).casefold()
results=[]
for step,fn in CHECKPOINTS:
    p=d/fn
    if not p.is_file(): print(f"STEP {step}: MISSING {fn}",flush=True); continue
    print(f"\n=== STEP-{step} UNSEEN EXACT-MATCH EVAL ===",flush=True)
    model=mod.ScugnizzDecoder(cfg).to(device)
    w=torch.load(p,map_location=device,weights_only=True)
    model.load_state_dict(w["model"] if isinstance(w,dict) and "model" in w else w,strict=True); model.eval()
    passed=0
    for i,(name,msg,expected) in enumerate(TESTS,1):
        ans=generate(model,tok,prompt(msg),64,0.0,0,device).strip(); ok=norm(ans)==norm(expected); passed+=int(ok)
        print(f"[{i}/6] {name}\nEXPECTED: {expected!r}\nANSWER:   {ans!r}\nRESULT:   {'PASS' if ok else 'FAIL'}",flush=True)
    results.append((step,passed)); print(f"SUMMARY STEP {step}: {passed}/6 ({passed/6:.1%})",flush=True)
    del model
    if torch.cuda.is_available(): torch.cuda.empty_cache()
print("\n=== FINAL COMPARISON ===",flush=True)
for s,p in results: print(f"STEP {s}: {p}/6 ({p/6:.1%})",flush=True)
if results:
    best=max(p for s,p in results); print(f"BEST CHECKPOINT: {', '.join(s for s,p in results if p==best)} ({best}/6 = {best/6:.1%})",flush=True)