"""Run the direct H3 core against one matched ComfyUI capture."""

import argparse
import time

import torch

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


parser = argparse.ArgumentParser()
parser.add_argument("--capture-dir", default="/artifacts/capture")
parser.add_argument("--model", default="/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")
args = parser.parse_args()

capture_dir = args.capture_dir
inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
expected = torch.load(f"{capture_dir}/output.pt", map_location="cuda", weights_only=False)
checkpoint = H3Checkpoint(args.model)

start = time.perf_counter()
model = H3PackedDenoiser.from_checkpoint(checkpoint).eval()
with torch.inference_mode():
    video, audio = model(
        inputs["hidden"],
        inputs["timesteps"],
        inputs["position_ids"],
        inputs["segments"],
        expected["video_segment"],
        expected["audio_segment"],
    )
torch.cuda.synchronize()

for name, actual, reference in (("video", video, expected["video"]), ("audio", audio, expected["audio"])):
    reference = reference.reshape_as(actual)
    delta = (actual.float() - reference.float()).abs()
    print(f"{name} shape={tuple(actual.shape)} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
print(f"elapsed_s={time.perf_counter() - start:.3f}")