import sys, torch
from transformers import AutoModel, AutoTokenizer
from cases import CASES
import pack, q64ctx
from eos_graph import load_chunked, eos_inputs, with_slots, chunk_inputs
S, P = int(sys.argv[1]), int(sys.argv[2])
ref = AutoModel.from_pretrained("eos", trust_remote_code=True, device="mps")
tok = AutoTokenizer.from_pretrained("eos"); pad = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
g = load_chunked("eos", S, P); worst = 0
cases = CASES + [(q64ctx.state, dict(list(q64ctx.Q.items())[:9]))]
for state, qs in cases:
    r = ref.system_one(state=state, questions=qs)["answers"]
    jobs = pack.rows(tok, state, qs)
    x, cand, qry, owner, used = eos_inputs(g.cfg, [e for _, _, e in jobs], S, P, pad)
    x = with_slots(x, cand, qry, 64)
    x["seg_chunks"], x["cont"], x["last_seg"] = chunk_inputs(x["segment"])
    names = ("input_ids", "cos", "sin", "valid", "tail_onehot", "segment", "lag_keep", "lag_tail", "seg_chunks", "cont", "last_seg", "cand_idx", "query_idx")
    with torch.no_grad():
        lg = g(*[torch.from_numpy(x[k]) for k in names])
    a = pack.answers(jobs, lg.tolist(), owner, None)
    d = 0
    for q in qs:
        pr = r[q].get("probabilities", {"true": r[q].get("noul")}); pa = a[q].get("probabilities", {"true": a[q].get("noul")})
        d = max(d, max(abs(pr[k] - pa[k]) for k in pr))
    worst = max(worst, d); print(f"{len(qs)} questions, packed {used}, max dp {d:.2e}")
print("WORST", worst)