"""Identify the first direct/Comfy NVFP4 MLP divergence."""

import argparse

import torch

from h3_blackwell_runtime.block import gate_segments, modulate_segments
from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
from h3_blackwell_runtime.attention import rms_norm
from h3_blackwell_runtime.rope import h3_rope_rotation


parser = argparse.ArgumentParser()
parser.add_argument("--capture-dir", required=True)
parser.add_argument("--model", required=True)
parser.add_argument("--reference-input", action="store_true")
args = parser.parse_args()

inputs = torch.load(f"{args.capture_dir}/input.pt", map_location="cuda", weights_only=False)
model = H3PackedDenoiser.from_checkpoint(H3Checkpoint(args.model)).eval()
block = model.backbone.blocks[0]
adaln = model.backbone.adaln[0]
rotation = h3_rope_rotation(inputs["position_ids"], model.backbone.inv_freq, inputs["hidden"].dtype)
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = adaln(inputs["timesteps"])

with torch.inference_mode():
    norm1 = modulate_segments(rms_norm(inputs["hidden"], block.norm1_weight, block.norm_eps), shift_msa, scale_msa, inputs["segments"])
    post_attention = gate_segments(inputs["hidden"], block.attention(norm1, rotation), gate_msa, inputs["segments"])
    norm2 = modulate_segments(rms_norm(post_attention, block.norm2_weight, block.norm_eps), shift_mlp, scale_mlp, inputs["segments"])
    if args.reference_input:
        norm2 = torch.load(f"{args.capture_dir}/block0_norm2.pt", map_location="cuda", weights_only=False)
    fc1 = block.mlp.fc1(norm2)
    gate, up = fc1.chunk(2, dim=-1)
    activated = torch.nn.functional.silu(gate).mul_(up)
    fc2 = block.mlp.fc2(activated)

for name, actual in (("fc1", fc1), ("activated", activated), ("fc2", fc2)):
    expected = torch.load(f"{args.capture_dir}/block0_mlp_{name}.pt", map_location="cuda", weights_only=False)
    delta = (actual.float() - expected.float()).abs()
    print(f"{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")