"""Replay a matched Comfy keyframe sampler capture with exact static conditioning."""

import argparse
from pathlib import Path

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, beta_sigmas, res_multistep_update
from h3_blackwell_runtime.t2v import random_av_latents


parser = argparse.ArgumentParser()
parser.add_argument("--sampler", type=Path, required=True)
parser.add_argument("--dit", type=Path, required=True)
parser.add_argument("--model", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
parser.add_argument("--result-latent", type=Path)
parser.add_argument("--production-trace", type=Path)
args = parser.parse_args()

checkpoint = H3Checkpoint(args.model)
model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend="sage2").eval()
packer = H3PromptPacker(checkpoint)
captured_input = torch.load(args.dit / "input.pt", map_location="cuda", weights_only=False)
text_length = next(start for start, _, code in captured_input["segments"] if code == 3)
prefix_stop = next(start for start, _, code in captured_input["segments"] if code % 3 == 2)
text = captured_input["hidden"][:text_length].unsqueeze(0).cuda()
prefix = captured_input["hidden"][:prefix_stop].cuda()
text_tags = torch.ones(text_length, dtype=torch.long, device="cuda")
for start, stop, code in captured_input["segments"]:
    if stop <= text_length:
        text_tags[start:stop] = code % 3

video, audio, frame_count = random_av_latents(384, 384, 22, 440207)
video_shape, audio_shape = video.shape, audio.shape
video_count, audio_count = video.numel(), audio.numel()
zero_cond = [torch.zeros(1, 24, 1, 24, 24, device="cuda") for _ in range(2)]
sigmas = torch.load(args.sampler / "initial.pt", map_location="cuda", weights_only=False)["sigmas"].cuda()
direct_sigmas = beta_sigmas(12, device="cuda")
sigma_delta = (direct_sigmas.float() - sigmas.float()).abs()
print({"stage": "sigmas", "direct": direct_sigmas.tolist(), "comfy": sigmas.tolist(), "mean_delta": float(sigma_delta.mean()), "max_delta": float(sigma_delta.max())}, flush=True)
video_history = audio_history = history_sigma = None

for index, sigma in enumerate(sigmas[:-1]):
    reference = torch.load(args.sampler / f"step_{index:02d}.pt", map_location="cuda", weights_only=False)
    reference_x = reference["x"].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.float() - reference_audio.float()).abs()

    sigma_audio = _audio_sigma(sigma)
    carry = sigma_audio / sigma
    hidden, times, segments, positions, video_segment, audio_segment = packer(
        text,
        video,
        audio.to(torch.bfloat16) * carry,
        _model_sigma(sigma),
        text_token_tags=text_tags,
        cond_latents=zero_cond,
        cond_frame_indices=[0, frame_count - 1],
        frame_count=frame_count,
        seed=440207,
    )
    hidden[:prefix_stop] = prefix.to(hidden)
    input_hidden = hidden.detach().clone()
    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.to(torch.bfloat16) * carry.to(torch.bfloat16))
        + (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-_unpack_audio(raw_audio))
    ).float()
    audio_denoised = audio - sigma * audio_model_output

    reference_denoised = reference["denoised"].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()
    production = torch.load(args.production_trace / f"step_{index:02d}.pt", map_location="cuda", weights_only=False) if args.production_trace else None
    production_video = (production["video"].float() - reference_video.float()).abs() if production else None
    production_audio = (production["audio"].float() - reference_audio.float()).abs() if production else None
    production_denoised = (production["video_denoised"].float() - reference_video_denoised.float()).abs() if production else None
    production_hidden = (production["hidden"].float() - input_hidden.float()).abs() if production and "hidden" in production else None
    production_raw = (production["raw_video"].float() - raw_video.float()).abs() if production and "raw_video" in production else None
    production_segments = [
        (start, stop, code, float((production["hidden"][start:stop].float() - input_hidden[start:stop].float()).abs().mean()))
        for start, stop, code in segments
    ] if production and "hidden" in production else None
    print({
        "step": index,
        "pre_video_mean": float(pre_video.mean()),
        "pre_video_max": float(pre_video.max()),
        "pre_audio_mean": float(pre_audio.mean()),
        "pre_audio_max": float(pre_audio.max()),
        "denoised_video_mean": float(denoised_video.mean()),
        "denoised_video_max": float(denoised_video.max()),
        "denoised_audio_mean": float(denoised_audio.mean()),
        "denoised_audio_max": float(denoised_audio.max()),
        "production_video_mean": float(production_video.mean()) if production_video is not None else None,
        "production_video_max": float(production_video.max()) if production_video is not None else None,
        "production_audio_mean": float(production_audio.mean()) if production_audio is not None else None,
        "production_audio_max": float(production_audio.max()) if production_audio is not None else None,
        "production_denoised_mean": float(production_denoised.mean()) if production_denoised is not None else None,
        "production_denoised_max": float(production_denoised.max()) if production_denoised is not None else None,
        "production_hidden_mean": float(production_hidden.mean()) if production_hidden is not None else None,
        "production_hidden_max": float(production_hidden.max()) if production_hidden is not None else None,
        "production_raw_mean": float(production_raw.mean()) if production_raw is not None else None,
        "production_raw_max": float(production_raw.max()) if production_raw is not None else None,
        "production_segments": production_segments,
    }, flush=True)

    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 = res_multistep_update(audio, audio_denoised, sigma, sigma_down, audio_history, history_sigma, previous_sigma)
    video_history, audio_history, history_sigma = video_denoised, audio_denoised, sigma_down

if args.result_latent:
    result = torch.load(args.result_latent, map_location="cuda", weights_only=False)
    result_video = result["latent"] if isinstance(result, dict) else result
    delta = (video.float() - result_video.cuda().float()).abs()
    print({"stage": "final_video", "mean_delta": float(delta.mean()), "max_delta": float(delta.max())}, flush=True)