"""Compare direct FP32 patch GEMMs under Comfy's captured matmul policy."""

import argparse
from pathlib import Path

import torch
import torch.nn.functional as functional

from h3_blackwell_runtime.checkpoint import H3Checkpoint


parser = argparse.ArgumentParser()
parser.add_argument("--capture", type=Path, required=True)
parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
args = parser.parse_args()

captured = torch.load(args.capture, map_location="cuda", weights_only=False)
checkpoint = H3Checkpoint(args.checkpoint)
previous_tf32 = torch.backends.cuda.matmul.allow_tf32
previous_precision = torch.get_float32_matmul_precision()
torch.backends.cuda.matmul.allow_tf32 = captured["matmul"]["allow_tf32"]
torch.set_float32_matmul_precision(captured["matmul"]["precision"])
try:
    for name in ("video", "audio"):
        rows = captured[f"{name}_rows"].to("cuda")
        weight = checkpoint.tensor(f"{name}_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
        bias = checkpoint.tensor(f"{name}_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
        output = functional.linear(rows, weight, bias)
        for label, actual, expected in (("weight", weight, captured[f"{name}_weight"]), ("bias", bias, captured[f"{name}_bias"]), ("fp32", output, captured[f"{name}_embed_fp32"])):
            delta = (actual.float() - expected.to(actual.device).float()).abs()
            print(f"{name}.{label} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
        bf16_delta = (output.to(torch.bfloat16).float() - captured[f"{name}_embed_fp32"].to(torch.bfloat16).float()).abs()
        print(f"{name}.bf16 max_abs={bf16_delta.max().item():.6g} mean_abs={bf16_delta.mean().item():.6g}")
    print(f"matmul={captured['matmul']}")
finally:
    torch.backends.cuda.matmul.allow_tf32 = previous_tf32
    torch.set_float32_matmul_precision(previous_precision)