"""Compare direct and upstream MiniMax H3 VAE encoding on one keyframe."""

import argparse
import os
import sys

import numpy as np
from PIL import Image
import torch

from h3_blackwell_runtime.vae_encoder import (
    MiniMaxH3VideoVAEEncoder as DirectVAEEncoder,
    _causal_conv3d,
    _downsample,
    _group_norm_3d,
    _resnet,
)


parser = argparse.ArgumentParser()
parser.add_argument("--image", required=True)
parser.add_argument("--vae", default="/vae/minimax_h3_video_vae_fp16.safetensors")
parser.add_argument("--comfy-path", default="/opt/ComfyUI")
parser.add_argument("--tiling", action="store_true")
args = parser.parse_args()

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


def load_checkpoint(path):
    if os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}:
        from safetensors.torch import load

        with open(path, "rb") as file:
            return load(file.read())
    from safetensors.torch import load_file

    return load_file(path, device="cuda")


def report(name, actual, expected):
    actual = actual.detach()
    expected = expected.detach().to(actual.device)
    delta = (actual.float() - expected.float()).abs()
    print({
        "stage": name,
        "shape": tuple(actual.shape),
        "actual_min": float(actual.min()),
        "actual_max": float(actual.max()),
        "expected_min": float(expected.min()),
        "expected_max": float(expected.max()),
        "mean_delta": float(delta.mean()),
        "max_delta": float(delta.max()),
    }, flush=True)


image = Image.open(args.image).convert("RGB")
pixels = torch.from_numpy(np.asarray(image).copy()).permute(2, 0, 1).unsqueeze(0).cuda().float()
pixels = pixels.div(127.5).sub(1.0)

direct = DirectVAEEncoder.from_safetensors(args.vae, device="cuda", tiling=args.tiling).eval()
upstream = UpstreamVAE(tiling=args.tiling).to("cuda").eval()
upstream = upstream.to(dtype=direct.compute_dtype)
state = load_checkpoint(args.vae)
upstream.load_state_dict(state, strict=True)
del state

trace = {}
trace_names = {"encoder.conv_in", "encoder.norm_out", "encoder.conv_out", "quant_conv"}
for level in range(6):
    trace_names.update(f"encoder.down.{level}.block.{block}" for block in range(2))
    if level < 4:
        trace_names.add(f"encoder.down.{level}.downsample")
for name, module in upstream.named_modules():
    if name in trace_names:
        module.register_forward_hook(lambda _module, _inputs, output, name=name: trace.__setitem__(name, output.detach().cpu()))

report("quant_conv_weight", direct.quant_conv.weight, upstream.quant_conv.weight)
report("quant_conv_bias", direct.quant_conv.bias, upstream.quant_conv.bias)
with torch.inference_mode():
    direct_latent = direct.encode(pixels.clone())
    upstream_latent = upstream.encode(pixels.to(direct.compute_dtype))
report("normalized_latent", direct_latent, upstream_latent)

with torch.inference_mode():
    x = pixels.unsqueeze(2)
    x = (x + 1.0) * 0.5
    x = (x - direct.pixel_mean.to(x)) / direct.pixel_std.to(x)
    params = direct.W
    x = _causal_conv3d(x, params["conv_in"][0], params["conv_in"][1], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
    report("encoder.conv_in", x, trace.pop("encoder.conv_in"))
    for level_index, level in enumerate(params["down"]):
        for block_index, block in enumerate(level["blocks"]):
            x = _resnet(x, block)
            name = f"encoder.down.{level_index}.block.{block_index}"
            report(name, x, trace.pop(name))
        if level["down"] is not None:
            x = _downsample(x, level["down"])
            name = f"encoder.down.{level_index}.downsample"
            report(name, x, trace.pop(name))
    x = _group_norm_3d(x, params["norm_out_w"], params["norm_out_b"])
    report("encoder.norm_out", x, trace.pop("encoder.norm_out"))
    x = torch.nn.functional.silu(x)
    x = _causal_conv3d(x, params["conv_out"][0], params["conv_out"][1], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
    report("encoder.conv_out", x, trace.pop("encoder.conv_out"))
    x = torch.nn.functional.conv3d(x, direct.quant_conv.weight, direct.quant_conv.bias)
    report("quant_conv", x, trace.pop("quant_conv"))
    report("latents_mean", direct.latents_mean, upstream.latents_mean)
    report("latents_std", direct.latents_std, upstream.latents_std)