"""Compare direct FL2VA beta/RES steps with captured Comfy sampler state."""

import glob

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, _unpack_audio, res_multistep_update


root = "/artifacts/fl2va-sampler-reference"
initial = torch.load(f"{root}/initial.pt", map_location="cuda", weights_only=False)
steps = [torch.load(path, map_location="cuda", weights_only=False) for path in sorted(glob.glob(f"{root}/step_*.pt"))]
sigmas = initial["sigmas"].to("cuda")
checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
captured_input = torch.load("/artifacts/capture/input.pt", map_location="cuda", weights_only=False)
text = captured_input["hidden"][:17].unsqueeze(0)
model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend="sage2").eval()
packer = H3PromptPacker(checkpoint)

video_shape = (1, 24, 7, 12, 20)
audio_shape = (1, 32, 2, 37)
video_count = torch.tensor(video_shape).prod().item()
old_video = old_audio = old_sigma = None
for index, reference in enumerate(steps):
    sigma, sigma_down = sigmas[index], sigmas[index + 1]
    packed_state = reference["x"].to("cuda").reshape(-1)
    video = packed_state[:video_count].reshape(video_shape)
    audio_carried = packed_state[video_count:].reshape(audio_shape)
    native_audio = audio_carried * (_audio_sigma(sigma) / sigma)
    hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, float(sigma))
    raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
    velocity_video = -unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1])
    velocity_audio = -_unpack_audio(raw_audio)
    carry = _audio_sigma(sigma) / sigma
    velocity_audio = (1.0 - 4.0) * (audio_carried * carry) + (1.0 + 3.0 * _audio_sigma(sigma)) * velocity_audio
    denoised = (video - sigma * velocity_video, audio_carried - sigma * velocity_audio)
    reference_denoised = reference["denoised"].to("cuda").reshape(-1)
    reference_denoised_video = reference_denoised[:video_count].reshape(video_shape)
    denoised_delta = (denoised[0].float() - reference_denoised_video.float()).abs()
    previous_sigma = sigmas[index - 1] if index else None
    video = res_multistep_update(video, denoised[0], sigma, sigma_down, old_video, old_sigma, previous_sigma)
    audio_carried = res_multistep_update(audio_carried, denoised[1], sigma, sigma_down, old_audio, old_sigma, previous_sigma)
    if index + 1 < len(steps):
        reference_latent = steps[index + 1]["x"].to("cuda").reshape(-1)[:video_count].reshape(video_shape)
        latent_delta = (video.float() - reference_latent.float()).abs()
        latent_text = f"latent_video_mean={latent_delta.mean().item():.6g} latent_video_max={latent_delta.max().item():.6g}"
    else:
        latent_text = "latent_video_mean=final-unobserved latent_video_max=final-unobserved"
    print(f"step={index:02d} x0_video_mean={denoised_delta.mean().item():.6g} x0_video_max={denoised_delta.max().item():.6g} {latent_text}")
    old_video, old_audio, old_sigma = denoised[0], denoised[1], sigma_down