"""Trace direct/upstream temporal VAE assembly from matching tiled clips."""

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("--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 diff_stats(label: str, a: torch.Tensor, b: torch.Tensor) -> None:
    diff = (a.float() - b.float()).abs()
    print({"label": label, "shape": tuple(a.shape), "max": float(diff.max()), "mean": float(diff.mean())}, flush=True)


def plan(vae, z_len: int):
    pseudo = z_len + vae.token_drop
    pad = (-pseudo) % vae.tokens_chunk_size
    pseudo += pad
    chunks = pseudo // vae.tokens_chunk_size - int(vae.token_drop > 0)
    if chunks < 1:
        pad += vae.tokens_chunk_size
        chunks += 1
    return chunks, pad, vae._decode_temporal_frame_plan(z_len + pad, chunks, pad)


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=True).eval()
upstream = UpstreamVAE(tiling=True).to("cuda").eval()
upstream.load_state_dict(load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda"), strict=True)

with torch.inference_mode():
    zd = latent.to(next(direct.parameters()).dtype)
    zu = latent.to(next(upstream.parameters()).dtype)
    zd = zd * direct.latents_std.view(1, -1, 1, 1, 1).to(zd) + direct.latents_mean.view(1, -1, 1, 1, 1).to(zd)
    zu = zu * upstream.latents_std.view(1, -1, 1, 1, 1).to(zu) + upstream.latents_mean.view(1, -1, 1, 1, 1).to(zu)
    chunks, pad, frames = plan(direct, zd.shape[2])
    if pad:
        zd = torch.cat((zd, zd[:, :, -1:].repeat(1, 1, pad, 1, 1)), dim=2)
        zu = torch.cat((zu, zu[:, :, -1:].repeat(1, 1, pad, 1, 1)), dim=2)
    print({"chunks": chunks, "pad": pad, "frames": frames}, flush=True)

    chunk_dec = direct.tokens_chunk_size * direct.vae_ratio_t
    split_count = int(direct.token_drop > 0) + 1
    prev_d = prev_u = None
    write_pos = 0
    for chunk_index in range(chunks):
        start = chunk_index * direct.tokens_chunk_size
        end = start + direct.tokens_chunk_size + direct.token_overlap
        clip_d = direct.tiled_decode(zd[:, :, start:end])
        clip_u = upstream.tiled_decode(zu[:, :, start:end])
        diff_stats(f"clip_{chunk_index}", clip_d, clip_u)
        for split in range(split_count):
            frame_start = split * chunk_dec
            frame_end = min(frame_start + chunk_dec, clip_d.shape[2])
            part_d = clip_d[:, :, frame_start:frame_end][:, :, direct.frame_pre_padding:]
            part_u = clip_u[:, :, frame_start:frame_end][:, :, upstream.frame_pre_padding:]
            if split == 0 and prev_d is not None:
                diff_stats(f"chunk_{chunk_index}_preblend_current", part_d, part_u)
                diff_stats(f"chunk_{chunk_index}_direct_blend_same_inputs", direct.blend(prev_d, part_d, direct.frame_overlap, -3), upstream.blend(prev_d, part_d, direct.frame_overlap, -3))
                diff_stats(f"chunk_{chunk_index}_upstream_blend_same_inputs", direct.blend(prev_u, part_u, upstream.frame_overlap, -3), upstream.blend(prev_u, part_u, upstream.frame_overlap, -3))
                part_d = direct.blend(prev_d, part_d, direct.frame_overlap, -3)
                part_u = upstream.blend(prev_u, part_u, upstream.frame_overlap, -3)
                diff_stats(f"chunk_{chunk_index}_blended_at_{write_pos}", part_d, part_u)
            elif split == 0:
                diff_stats(f"chunk_{chunk_index}_write_at_{write_pos}", part_d, part_u)
            else:
                prev_d, prev_u = part_d.contiguous(), part_u.contiguous()
                diff_stats(f"chunk_{chunk_index}_overlap", prev_d, prev_u)
                continue
            write_pos += part_d.shape[2]