"""Compare direct and upstream H3 VAE decoder internals on one latent clip."""

import argparse
import sys
from pathlib import Path

import torch
from safetensors.torch import load_file

from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE as DirectVAE


parser = argparse.ArgumentParser()
parser.add_argument("--latent", type=Path, required=True)
parser.add_argument("--start", type=int, default=0)
parser.add_argument("--tokens", type=int, default=8)
parser.add_argument("--comfy-path", default="/opt/ComfyUI")
args = parser.parse_args()

sys.path.insert(0, args.comfy_path)
from h3_blackwell_runtime.upstream_vae import MiniMaxH3VideoVAE as UpstreamVAE  # noqa: E402


def stats(name: str, a: torch.Tensor, b: torch.Tensor) -> None:
    diff = (a.float() - b.float()).abs()
    print({"stage": name, "max": float(diff.max()), "mean": float(diff.mean()), "shape": tuple(a.shape)}, flush=True)


state = torch.load(args.latent, map_location="cuda", weights_only=False)
latent = state["latent"].to("cuda") if isinstance(state, dict) else state.to("cuda")

direct = DirectVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda", tiling=False).eval()
upstream = UpstreamVAE(tiling=False).to("cuda").eval()
upstream.load_state_dict(load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda"), strict=True)

with torch.inference_mode():
    z_d = latent[:, :, args.start:args.start + args.tokens].to(next(direct.parameters()).dtype)
    z_u = z_d.clone().to(next(upstream.parameters()).dtype)
    z_d = z_d * direct.latents_std.view(1, -1, 1, 1, 1).to(z_d) + direct.latents_mean.view(1, -1, 1, 1, 1).to(z_d)
    z_u = z_u * upstream.latents_std.view(1, -1, 1, 1, 1).to(z_u) + upstream.latents_mean.view(1, -1, 1, 1, 1).to(z_u)
    z_d = direct.post_quant_conv(z_d)
    z_u = upstream.post_quant_conv(z_u)
    stats("post_quant_conv", z_d, z_u)

    dd, du = direct.decoder, upstream.decoder
    h_d = dd.x_embedder(z_d.flatten(2).transpose(1, 2))
    h_u = du.x_embedder(z_u.flatten(2).transpose(1, 2))
    stats("x_embedder", h_d, h_u)

    b, _, latent_t, latent_h, latent_w = z_d.shape
    h_d = torch.cat((h_d, dd.register_tokens.to(h_d).expand(b, -1, -1), torch.zeros_like(h_d[:, :1])), dim=1)
    h_u = torch.cat((h_u, upstream.decoder.register_tokens.to(h_u).expand(b, -1, -1), torch.zeros_like(h_u[:, :1])), dim=1)
    ids_d = __import__("h3_blackwell_runtime.vae_decoder", fromlist=["create_token_ids"]).create_token_ids((latent_t, latent_h, latent_w), z_d.device, z_d.dtype).expand(b, -1, -1)
    ids_d = torch.cat((ids_d, torch.zeros(b, 1 + dd.num_register_tokens, 3, device=z_d.device, dtype=z_d.dtype)), dim=1)
    ids_u = __import__("h3_blackwell_runtime.upstream_vae", fromlist=["create_token_ids"]).create_token_ids((latent_t, latent_h, latent_w), z_u.device, z_u.dtype).expand(b, -1, -1)
    ids_u = torch.cat((ids_u, torch.zeros(b, 1 + du.num_register_tokens, 3, device=z_u.device, dtype=z_u.dtype)), dim=1)
    rope_d = dd.pos_embed(ids_d)
    rope_u = du.pos_embed(ids_u)
    stats("rope", rope_d, rope_u)

    for index, (block_d, block_u) in enumerate(zip(dd.transformer_blocks, du.transformer_blocks)):
        h_d = block_d(h_d, rope_d)
        h_u = block_u(h_u, rope_u)
        stats(f"block_{index:02d}", h_d, h_u)

    patches = h_d.shape[1] - 1 - dd.num_register_tokens
    out_d = dd.proj_out(dd.norm_out(h_d))[:, :patches]
    out_u = du.proj_out(du.norm_out(h_u))[:, :patches]
    stats("proj_out", out_d, out_u)

    out_d = out_d.view(b, latent_t, latent_h, latent_w, dd.out_channels, dd.patch_size_t, dd.patch_size, dd.patch_size)
    out_u = out_u.view(b, latent_t, latent_h, latent_w, du.out_channels, du.patch_size_t, du.patch_size, du.patch_size)
    out_d = out_d.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(b, dd.out_channels, latent_t * dd.patch_size_t, latent_h * dd.patch_size, latent_w * dd.patch_size)
    out_u = out_u.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(b, du.out_channels, latent_t * du.patch_size_t, latent_h * du.patch_size, latent_w * du.patch_size)
    stats("decoded", out_d, out_u)

    pix_d = out_d.float().mul_(direct.pixel_std.to(out_d)).add_(direct.pixel_mean.to(out_d)).clamp_(0, 1).mul_(2).sub_(1)
    pix_u = out_u.float().mul_(upstream.pixel_std.to(out_u)).add_(upstream.pixel_mean.to(out_u)).clamp_(0, 1).mul_(2).sub_(1)
    stats("pixels", pix_d, pix_u)