import sys, json, random, collections, torch
from transformers import AutoModelForCausalLM, AutoConfig
from stokenizer import STokenizer
from coconut import Coconut

CKPT = sys.argv[1]
DEPTH = 1
N = 40

tok = STokenizer()
lat = tok.convert_tokens_to_ids("<|latent|>")
s0 = tok.convert_tokens_to_ids("<|start-latent|>")
e0 = tok.convert_tokens_to_ids("<|end-latent|>")
m = AutoModelForCausalLM.from_config(AutoConfig.from_pretrained("configs/symbol-2layer-8head-768dim.json"))
m = Coconut(m, lat, s0, e0, tok.eos_token_id)
sd = torch.load(CKPT, map_location="cpu")
sd = {k[len("module."):] if k.startswith("module.") else k: v for k, v in sd.items()}
print("load:", m.load_state_dict(sd, strict=False))
m.eval()

data = json.load(open("data/star_2arm_L6_valid_fo_coconut.json"))
cnt = collections.Counter()
shown = 0
for i, s in enumerate(data[:N]):
    reach = s["neighbor_k"][str(DEPTH)][0]
    neg = s["neg_neighbor_k"][str(DEPTH)][0]
    cands = [reach, neg]
    random.Random(i).shuffle(cands)              # same seeding as finalonly_categorize
    reach_pos = cands.index(reach)
    q = ("<eos> " + "|".join([f" {e[0]} {e[1]} " for e in s["edges"]]).strip()
         + " [Q] " + str(cands[0]) + " " + str(cands[1])
         + " [R] " + str(s["root"]) + " <|latent|>" * DEPTH + " [A] ")
    ids = torch.tensor([tok.encode(q, add_special_tokens=False)])
    with torch.no_grad():
        out = m.generate(ids, torch.ones_like(ids), max_new_tokens=1, synced_gpus=False)
    g = int(out[0, -1].item())
    where = "cand0" if g == cands[0] else ("cand1" if g == cands[1] else "OTHER")
    cnt["reach_pos%d" % reach_pos] += 1
    cnt[where] += 1
    cnt["correct" if g == reach else "wrong"] += 1
    if shown < 16:
        print(f"root={s['root']:2d} [Q]={cands} reach={reach}(pos{reach_pos}) neg={neg} | out={g}[{where}] {'OK' if g==reach else 'x'}")
        shown += 1
print("\nTALLY:", dict(cnt))