#!/usr/bin/env python3
"""Check every real checkpoint key/shape against native Comfy on the meta device."""
import argparse
import json
from pathlib import Path

import torch

import convert_omniref_comfy as conv
from test_conversion import native_namespace


def audit(src):
    c, source = conv.load_source(src)
    expected = conv.expected_shapes(c)
    assert set(source) == set(expected)
    assert all(source[k].shape == shape for k, shape in expected.items())
    native = native_namespace()
    receipts = []
    for rank in [0, 8]:
        target = {}
        for name, tensor in source.items():
            if ".attn.to_k." in name or ".attn.to_v." in name:
                continue
            shape = list(tensor.shape)
            output_name = conv.renamed(name)
            if ".attn.to_q." in name:
                output_name = output_name.replace(".attn.to_q.", ".attn.qkv_proj.")
                shape[0] *= 3
            if rank and name.startswith("time_embedder."):
                continue
            if rank and (name == "norm_out.linear.weight" or name.endswith(".adaln_proj.linear.weight")):
                shape[-1] = rank
            target[output_name] = torch.empty(shape, device="meta")
        target["rope.inv_freq"] = torch.empty([c["rope_freq_dim"]], device="meta")
        if rank:
            target["adaln_t_table"] = torch.empty([1025, rank], device="meta")
        metadata = conv.metadata_for(c, {"video_scheduler_shift": 12.0, "audio_scheduler_shift": 3.0}, rank)
        detected = native["detect_unet_config"](target, "", metadata=metadata)
        model = native["MiniMaxH3Model"](**detected, dtype=torch.float32, device="meta", operations=torch.nn)
        result = model.load_state_dict(target, strict=True, assign=True)
        assert not result.missing_keys and not result.unexpected_keys
        assert all(target[k].shape == model.state_dict()[k].shape for k in target)
        receipts.append({"source_tensors": len(source), "target_tensors": len(target),
                         "curve_rank": rank, "detected_config": detected,
                         "native_strict_meta_load": "passed", "gpu_inference": "not_run"})
    return receipts


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--src", required=True, type=Path)
    args = parser.parse_args()
    print(json.dumps(audit(args.src), indent=2))