"""Replay the direct token refiner against a Comfy boundary trace."""

import argparse
from pathlib import Path

import torch
from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.token_refiner import H3TokenRefiner


parser = argparse.ArgumentParser()
parser.add_argument("--trace-dir", type=Path, required=True)
parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
args = parser.parse_args()

refiner = H3TokenRefiner(H3Checkpoint(args.checkpoint), attention_backend="sage2").eval()
hidden = torch.load(args.trace_dir / "refiner_input.pt", map_location="cuda", weights_only=False)

with torch.inference_mode():
    for index, block in enumerate(refiner.blocks):
        hidden = block(hidden)
        expected = torch.load(args.trace_dir / f"refiner_block{index}.pt", map_location="cuda", weights_only=False)
        delta = (hidden.float() - expected.float()).abs()
        print(f"block={index} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
    output = refiner.final_norm
    output = torch.nn.functional.rms_norm(hidden, output.shape, weight=output.to(hidden), eps=1e-5).unsqueeze(0)
    expected = torch.load(args.trace_dir / "refiner_output.pt", map_location="cuda", weights_only=False)
    delta = (output.float() - expected.float()).abs()
    print(f"output max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")