"""Compare direct prompt-only FL2VA packing with an immutable Comfy DiT input."""

import argparse
from pathlib import Path

import torch
import torch.nn.functional as functional

from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.packing import H3PromptPacker, pack_audio, patchify_video
from h3_blackwell_runtime.sampler import _audio_sigma


parser = argparse.ArgumentParser()
parser.add_argument("--capture", type=Path, required=True)
parser.add_argument("--refiner-trace", 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)
text = torch.load(args.refiner_trace / "refiner_output.pt", map_location="cuda", weights_only=False)
if text.ndim == 2:
    text = text.unsqueeze(0)
checkpoint = H3Checkpoint(args.checkpoint)
packer = H3PromptPacker(checkpoint)

video_shape = (1, 24, 7, 12, 20)
audio_shape = (1, 32, 2, 37)
video = captured["video_x"].to("cuda")
audio = captured["audio_x"].to("cuda")
sigma = 1 - captured["timesteps"].min()

with torch.inference_mode():
    hidden, times, segments, positions, _, _ = packer(text, video, audio, float(sigma))
    text_rows = text[0].to(torch.bfloat16)
    audio_rows = functional.linear(pack_audio(audio).float(), packer.audio_weight, packer.audio_bias).to(torch.bfloat16)
    video_rows = functional.linear(patchify_video(video).float(), packer.video_weight, packer.video_bias).to(torch.bfloat16)

for name, actual, expected in (
    ("text", text_rows, captured["hidden"][:17]),
    ("audio", audio_rows, captured["hidden"][17:91]),
    ("video", video_rows, captured["hidden"][91:]),
    ("hidden", hidden, captured["hidden"]),
    ("times", times, captured["timesteps"]),
    ("positions", positions, captured["position_ids"]),
):
    delta = (actual.float() - expected.to(actual.device).float()).abs()
    print(f"{name} shape={tuple(actual.shape)} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
print(f"segments direct={segments} comfy={captured['segments']}")