"""Compare direct NVFP4 wrapper and effective Comfy BF16 weights."""

import argparse

import torch
import torch.nn.functional as functional

from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder


parser = argparse.ArgumentParser()
parser.add_argument("--capture-dir", required=True)
parser.add_argument("--checkpoint", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors")
parser.add_argument("--dtype", choices=("bfloat16", "float32"), default="bfloat16")
args = parser.parse_args()

encoder = Qwen3VL32BTextEncoder(args.checkpoint, dtype=getattr(torch, args.dtype)).eval()
layer = encoder.layers[0]
qkv = torch.load(f"{args.capture_dir}/qwen0_quantized_qkv.pt", map_location="cuda", weights_only=False)
mlp = torch.load(f"{args.capture_dir}/qwen0_quantized_mlp.pt", map_location="cuda", weights_only=False)

for group, capture, modules, inputs in (
    ("qkv", qkv, {"q": layer.q_proj, "k": layer.k_proj, "v": layer.v_proj}, {"q": qkv["input"], "k": qkv["input"], "v": qkv["input"]}),
    ("mlp", mlp, {"gate": layer.gate_proj, "up": layer.up_proj, "down": layer.down_proj}, {"gate": mlp["input"], "up": mlp["input"], "down": mlp["activated"]}),
):
    for name, module in modules.items():
        expected = capture["modules"][name]
        packed = module._packed_weight()
        params = packed._params
        qdata, scale, block_scale = packed.layout_cls.get_plain_tensors(packed)
        reference_qdata = expected["qdata"].as_subclass(torch.Tensor).cpu()
        print(f"{group}.{name}.qdata exact={torch.equal(qdata.cpu(), reference_qdata)} shape={tuple(qdata.shape)}")
        for field, actual in (("scale", scale), ("block_scale", block_scale)):
            reference = expected[field].to(actual.device)
            print(f"{group}.{name}.{field} exact={torch.equal(actual, reference)} shape={tuple(actual.shape)} dtype={actual.dtype}")
        value = inputs[name].to(encoder.dtype)
        with torch.inference_mode():
            dequantized = packed.dequantize()
            direct = module(value)
            effective = expected["effective_weight"].to("cuda")
            comfy_effective = functional.linear(value.to(effective.dtype), effective, None)
        print(f"{group}.{name}.effective dtype={expected['effective_dtype']} shape={expected['effective_shape']} direct_dtype={encoder.dtype}")
        for variant, actual in (("weight", dequantized), ("linear", direct), ("effective_linear", comfy_effective)):
            reference = expected["effective_weight"].to(actual.device) if variant == "weight" else expected["output"].to(actual.device)
            delta = (actual.float() - reference.float()).abs()
            print(f"{group}.{name}.{variant}_vs_comfy max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")