"""Decode the captured final H3 video latent and compare against Comfy frames."""

import argparse
from pathlib import Path
import subprocess

import numpy as np
from PIL import Image
import torch

from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE


def load_reference_frames(path: Path, frames: int) -> torch.Tensor:
    files = sorted(path.glob("fl2va-reference_*.png"))[-frames:]
    if len(files) != frames:
        raise FileNotFoundError(f"expected {frames} reference PNGs in {path}, found {len(files)}")
    arrays = [np.asarray(Image.open(file).convert("RGB"), dtype=np.uint8) for file in files]
    return torch.from_numpy(np.stack(arrays, axis=0))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--latent", default="/artifacts/fl2va-sampler-reference/step_11.pt")
    parser.add_argument("--frames", type=int, default=22)
    parser.add_argument("--reference-frames", type=Path, default=Path("/artifacts/fl2va-sampler-reference-e679178/reference_frames"))
    parser.add_argument("--output", type=Path, default=Path("/output/h3-blackwell-runtime/direct-vae-same-latent.mp4"))
    args = parser.parse_args()

    state = torch.load(args.latent, map_location="cuda", weights_only=False)
    if "denoised" in state:
        video_shape = (1, 24, 7, 12, 20)
        video_count = torch.tensor(video_shape).prod().item()
        latent = state["denoised"].to("cuda").reshape(-1)[:video_count].reshape(video_shape)
    else:
        latent = state["video"].to("cuda")
    vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval()
    with torch.inference_mode():
        decoded = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, : args.frames]
    pixels = ((decoded[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu()
    reference = load_reference_frames(args.reference_frames, args.frames)
    delta = (pixels.to(torch.int16) - reference.to(torch.int16)).abs()

    args.output.parent.mkdir(parents=True, exist_ok=True)
    raw = args.output.with_suffix(".rgb")
    pixels.numpy().tofile(raw)
    subprocess.run(
        [
            "ffmpeg",
            "-y",
            "-f",
            "rawvideo",
            "-pixel_format",
            "rgb24",
            "-video_size",
            f"{pixels.shape[2]}x{pixels.shape[1]}",
            "-framerate",
            "24",
            "-i",
            str(raw),
            "-an",
            "-c:v",
            "libx264",
            "-pix_fmt",
            "yuv420p",
            str(args.output),
        ],
        check=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    raw.unlink()
    print(
        {
            "output": str(args.output),
            "shape": tuple(pixels.shape),
            "mean_abs_pixel_delta": float(delta.float().mean()),
            "max_abs_pixel_delta": int(delta.max()),
        }
    )


if __name__ == "__main__":
    main()