"""Heavier eval of v6 spec heads: multiple windows x more anchors."""
import sys
import torch

sys.path.insert(0, "/root/spec100m")
from config import Config
from model import build_model
from data_pipeline import load_tokens
from train_medusa import eval_streak

device = "cuda"
cfg = Config.v5_500m()
model = build_model(cfg, device)
data = load_tokens("mixture500m").to(device)

for ckpt_path in sys.argv[1:]:
    ckpt = torch.load(ckpt_path, map_location=device)
    missing, unexpected = model.load_state_dict(ckpt["model"], strict=False)
    model.eval()
    streaks, accs = [], []
    for _ in range(8):  # 8 windows x 64 anchors = 512 samples
        s, a = eval_streak(model, cfg, data, seq=1024, n_anchors=64, device=device)
        streaks.append(s)
        accs.append(a)
    sm = sum(streaks) / len(streaks)
    am = sum(accs) / len(accs)
    print(f"{ckpt_path} (step {ckpt.get('step')}): "
          f"streak {sm:.2f} (min {min(streaks):.2f} max {max(streaks):.2f})  "
          f"head-0 acc {am:.2f}")