"""Verify that TaoTrain checkpoint weights match the exported HF package exactly."""

from pathlib import Path

import torch

from export_to_hf import normalize_checkpoint, sanitize_model_state
from modeling_taonet import TaoNetForCausalLM


def summarize_tensor_diff(left: torch.Tensor, right: torch.Tensor) -> str:
    if left.shape != right.shape:
        return f"shape mismatch: {tuple(left.shape)} != {tuple(right.shape)}"
    if left.dtype != right.dtype:
        right = right.to(dtype=left.dtype)
    if left.dtype.is_floating_point:
        diff = (left - right).abs()
        return f"max_abs_diff={diff.max().item():.8g}, mean_abs_diff={diff.mean().item():.8g}"
    unequal = (left != right).sum().item()
    return f"unequal_values={unequal}"


def main():
    repo_dir = Path(__file__).resolve().parent
    checkpoint_path = repo_dir / "checkpoints" / "sft" / "final_model.pt"

    checkpoint = torch.load(checkpoint_path, map_location="cpu")
    checkpoint_state, _ = normalize_checkpoint(checkpoint)
    checkpoint_state, ignored_keys = sanitize_model_state(checkpoint_state)

    model = TaoNetForCausalLM.from_pretrained(str(repo_dir))
    exported_state = model.model.state_dict()

    checkpoint_keys = set(checkpoint_state)
    exported_keys = set(exported_state)

    missing = sorted(exported_keys - checkpoint_keys)
    unexpected = sorted(checkpoint_keys - exported_keys)

    print(f"checkpoint tensors: {len(checkpoint_keys)}")
    print(f"exported tensors:   {len(exported_keys)}")
    if ignored_keys:
        print(f"ignored checkpoint buffers: {len(ignored_keys)}")

    if missing:
        print("\nKeys present in exported model but missing from checkpoint:")
        for key in missing[:50]:
            print(f"  - {key}")

    if unexpected:
        print("\nKeys present in checkpoint but missing from exported model:")
        for key in unexpected[:50]:
            print(f"  - {key}")

    common_keys = sorted(checkpoint_keys & exported_keys)
    mismatches = []
    exact_matches = 0
    for key in common_keys:
        left = checkpoint_state[key].cpu()
        right = exported_state[key].cpu()
        if left.shape == right.shape and torch.equal(left, right.to(dtype=left.dtype) if right.dtype != left.dtype else right):
            exact_matches += 1
            continue
        mismatches.append((key, summarize_tensor_diff(left, right)))

    print(f"\nexact tensor matches: {exact_matches}/{len(common_keys)}")
    print(f"tensor mismatches:    {len(mismatches)}")

    if mismatches:
        print("\nFirst mismatches:")
        for key, summary in mismatches[:50]:
            print(f"  - {key}: {summary}")
        raise SystemExit(1)

    if missing or unexpected:
        raise SystemExit(1)

    print("\nWeight verification passed.")


if __name__ == "__main__":
    main()