"""Compare ComfyUI curve interpolation and block-0 AdaLN projection exactly."""

import torch
import torch.nn.functional as functional

from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.denoiser import H3PackedDenoiser


capture_dir = "/artifacts/capture"
inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
expected = torch.load(f"{capture_dir}/block0_norm1_adaln.pt", map_location="cuda", weights_only=False)
model = H3PackedDenoiser.from_checkpoint(H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")).eval()
adaln = model.backbone.adaln[0]
position = inputs["timesteps"].float().clamp(0, 1) * (adaln.curve_table.shape[0] - 1)
lower = position.floor().long().clamp(max=adaln.curve_table.shape[0] - 2)
embedding = torch.lerp(adaln.curve_table[lower], adaln.curve_table[lower + 1], (position - lower).unsqueeze(1))
values = functional.linear(embedding, adaln.weight, adaln.bias).view(embedding.shape[0] * 3, 6 * adaln.hidden_size)
shift, scale, *_ = values.chunk(6, dim=-1)

for name, actual, reference in (("embedding", embedding, inputs["t_emb"]), ("weight", adaln.weight, expected["weight"]), ("bias", adaln.bias, expected["bias"]), ("shift", shift, expected["shift"]), ("scale", scale, expected["scale"])):
    delta = (actual.float() - reference.float()).abs()
    print(f"{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")