"""Compare direct free-run FL2VA sampling against captured Comfy sampler state."""

import argparse

import torch

from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
from h3_blackwell_runtime.packing import H3PromptPacker, unpatchify_video
from h3_blackwell_runtime.sampler import _audio_sigma, _model_sigma, _unpack_audio, res_multistep_update
from h3_blackwell_runtime.t2v import random_av_latents


parser = argparse.ArgumentParser()
parser.add_argument("--root", default="/artifacts/fl2va-sampler-reference")
parser.add_argument("--capture", default="/artifacts/capture")
parser.add_argument("--width", type=int, default=320)
parser.add_argument("--height", type=int, default=192)
parser.add_argument("--frames", type=int, default=22)
parser.add_argument("--seed", type=int, default=440204)
parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2")
parser.add_argument("--oracle-timesteps", action="store_true")
args = parser.parse_args()

checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend=args.attention).eval()
packer = H3PromptPacker(checkpoint)
text = torch.load(f"{args.capture}/input_00.pt", map_location="cuda", weights_only=False)["hidden"][:17].unsqueeze(0).to("cuda")
video, audio_carried, _ = random_av_latents(args.width, args.height, args.frames, args.seed)
sigmas = torch.load(f"{args.root}/initial.pt", map_location="cuda", weights_only=False)["sigmas"].to("cuda")

video_shape = video.shape
audio_shape = audio_carried.shape
video_count = video.numel()
audio_count = audio_carried.numel()
video_history = audio_history = history_sigma = None

for index, sigma in enumerate(sigmas[:-1]):
    reference = torch.load(f"{args.root}/step_{index:02d}.pt", map_location="cuda", weights_only=False)
    h3_input = torch.load(f"{args.capture}/input_{index:02d}.pt", map_location="cuda", weights_only=False)
    reference_x = reference["x"].to("cuda").reshape(-1)
    reference_video = reference_x[:video_count].reshape(video_shape)
    reference_audio = reference_x[video_count:video_count + audio_count].reshape(audio_shape)
    pre_video = (video.float() - reference_video.float()).abs()
    pre_audio = (audio_carried.float() - reference_audio.float()).abs()

    sigma_audio = _audio_sigma(sigma)
    carry = sigma_audio / sigma
    model_timesteps = h3_input["timesteps"] if args.oracle_timesteps else None
    hidden, times, segments, positions, video_segment, audio_segment = packer(
        text,
        video,
        audio_carried.to(torch.bfloat16) * carry,
        _model_sigma(sigma),
        model_timesteps,
    )
    hidden_delta = (hidden.float() - h3_input["hidden"].to("cuda").float()).abs()
    time_delta = (times.float() - h3_input["timesteps"].to("cuda").float()).abs()

    with torch.inference_mode():
        raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
    raw_video = raw_video.to(torch.bfloat16).float()
    raw_audio = raw_audio.to(torch.bfloat16)
    video_denoised = video + sigma * unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1])
    audio_model_output = (
        (1.0 - 4.0) * (audio_carried.to(torch.bfloat16) * carry.to(torch.bfloat16))
        + (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-_unpack_audio(raw_audio))
    ).float()
    audio_denoised = audio_carried - sigma * audio_model_output

    reference_denoised = reference["denoised"].to("cuda").reshape(-1)
    reference_video_denoised = reference_denoised[:video_count].reshape(video_shape)
    reference_audio_denoised = reference_denoised[video_count:video_count + audio_count].reshape(audio_shape)
    denoised_video = (video_denoised.float() - reference_video_denoised.float()).abs()
    denoised_audio = (audio_denoised.float() - reference_audio_denoised.float()).abs()

    print(
        f"step={index:02d} "
        f"pre_max={max(pre_video.max().item(), pre_audio.max().item()):.6g} "
        f"hidden_max={hidden_delta.max().item():.6g} "
        f"time_max={time_delta.max().item():.6g} "
        f"denoised_max={max(denoised_video.max().item(), denoised_audio.max().item()):.6g} "
        f"denoised_mean={max(denoised_video.mean().item(), denoised_audio.mean().item()):.6g}"
    )

    previous_sigma = sigmas[index - 1] if index else None
    sigma_down = sigmas[index + 1]
    video = res_multistep_update(video, video_denoised, sigma, sigma_down, video_history, history_sigma, previous_sigma)
    audio_carried = res_multistep_update(audio_carried, audio_denoised, sigma, sigma_down, audio_history, history_sigma, previous_sigma)
    video_history, audio_history, history_sigma = video_denoised, audio_denoised, sigma_down