"""Verify the direct NVFP4 fc1 layout matches Comfy's loaded tensor."""

import argparse

import torch

from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.denoiser import H3PackedDenoiser


parser = argparse.ArgumentParser()
parser.add_argument("--capture-dir", required=True)
parser.add_argument("--model", required=True)
args = parser.parse_args()

expected = torch.load(f"{args.capture_dir}/block0_mlp_fc1_weight.pt", map_location="cuda", weights_only=False)
actual = H3PackedDenoiser.from_checkpoint(H3Checkpoint(args.model)).backbone.blocks[0].mlp.fc1
for name, value in (("qdata", actual.weight), ("scale", actual.weight_scale_2), ("block_scale", actual.weight_scale)):
    reference = expected[name]
    same = torch.equal(value, reference)
    delta = (value.float() - reference.float()).abs().max().item()
    print(f"{name} shape={tuple(value.shape)} dtype={value.dtype} exact={same} max_abs={delta:.6g}")
print(f"output_dtype={actual.output_dtype} reference_orig_dtype={expected['orig_dtype']}")
print(f"input_features={actual.in_features} output_features={actual.out_features} reference_orig_shape={expected['orig_shape']}")