"""Check that every tensor in a converted Hugging Face D-FINE checkpoint is bit-identical to a
tensor (or a q/k/v third of an in_proj tensor) of the official checkpoint.

  python verify_official.py Peterande/D-FINE 19214e91f5ce4c1e609022c38b1fc953c47404e7 dfine_s_coco.pth \
      ustc-community/dfine-small-coco f79e65b5fbb33ceb9d3ebba042955d7410c608f8
"""
import hashlib
import json
import sys

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

off_repo, off_rev, off_file, port_repo, port_rev = sys.argv[1:6]


def digest(t):
    t = t.detach().contiguous().cpu()
    return hashlib.sha256(str(t.dtype).encode() + str(tuple(t.shape)).encode() + t.numpy().tobytes()).hexdigest()


off_path = hf_hub_download(off_repo, off_file, revision=off_rev)
ckpt = torch.load(off_path, map_location="cpu", weights_only=False)
official = ckpt["model"]
pool = {}
for k, t in official.items():
    if not torch.is_tensor(t):
        continue
    pool.setdefault(digest(t), k)
    if "in_proj" in k and t.shape[0] % 3 == 0:
        for i, part in enumerate(t.chunk(3, dim=0)):
            pool.setdefault(digest(part), f"{k}[{'qkv'[i]}]")

port_path = hf_hub_download(port_repo, "model.safetensors", revision=port_rev)
port = load_file(port_path)
matched, unmatched = 0, []
for k, t in port.items():
    if digest(t) in pool:
        matched += 1
    else:
        unmatched.append((k, tuple(t.shape)))
nparams = sum(t.numel() for t in port.values())
unparams = sum(port[k].numel() for k, _ in unmatched)
print(json.dumps({
    "official": f"{off_repo}@{off_rev}/{off_file}",
    "official_sha256": hashlib.sha256(open(off_path, "rb").read()).hexdigest(),
    "official_keys": list(ckpt.keys()),
    "port": f"{port_repo}@{port_rev}/model.safetensors",
    "port_tensors": len(port),
    "matched": matched,
    "unmatched": len(unmatched),
    "unmatched_params": unparams,
    "total_params": nparams,
    "unmatched_sample": unmatched[:25],
}, indent=1))