"""Characterize isolated cuBLASLt NVFP4 scheduling at the real H3 FC2 boundary."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import platform
import statistics
import subprocess
import sys
import time
import types
from pathlib import Path
from typing import Any, Callable

import torch

from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
from h3_blackwell_runtime.nvfp4_quant import vortex_native_quantize_swiglu_nvfp4
from h3_blackwell_runtime.packing import H3PromptPacker
from h3_blackwell_runtime.rope import h3_rope_rotation
from h3_blackwell_runtime.sampler import _audio_sigma, _model_sigma, beta_sigmas, sample_video_res_multistep
from h3_blackwell_runtime.t2v import random_av_latents


ROOT = Path(__file__).resolve().parents[1]
SOURCE = ROOT / "research" / "fc2_nvfp4_scheduling" / "fc2_nvfp4_lt.cpp"
DEFAULT_BUDGETS = (0, 4 << 20, 8 << 20, 16 << 20, 32 << 20, 64 << 20)
DEFAULT_BLOCKS = (0, 24, 49)
ENV_PREFIXES = ("H3_", "COMFY_KITCHEN_", "CUDA_", "TORCH_")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--mode", choices=("compile", "characterize", "sweep", "selected", "profile", "block-gate", "trajectory"), default="characterize")
    parser.add_argument("--model-path", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
    parser.add_argument("--capture", type=Path, help="Directory containing input.pt, or an input .pt file")
    parser.add_argument("--output", type=Path, default=Path("/output/h3-blackwell-runtime/benchmarks/fc2-nvfp4-scheduling.json"))
    parser.add_argument("--build-directory", type=Path)
    parser.add_argument("--verbose-build", action="store_true")
    parser.add_argument("--width", type=int, default=1344)
    parser.add_argument("--height", type=int, default=768)
    parser.add_argument("--frames", type=int, default=124)
    parser.add_argument("--steps", type=int, default=12)
    parser.add_argument("--sampler-step", type=int, default=1)
    parser.add_argument("--seed", type=int, default=440420)
    parser.add_argument("--text-tokens", type=int, default=100)
    parser.add_argument("--blocks", type=int, nargs="+", default=list(DEFAULT_BLOCKS))
    parser.add_argument("--probe-block", type=int, default=24)
    parser.add_argument("--workspace-budgets", type=int, nargs="+", default=list(DEFAULT_BUDGETS))
    parser.add_argument("--requested-count", type=int, default=32)
    parser.add_argument("--explicit-split-k", type=int, nargs="+", default=[1, 2, 4, 8, 16])
    parser.add_argument("--max-explicit-checks", type=int, default=64)
    parser.add_argument("--candidate-limit", type=int, default=8)
    parser.add_argument("--candidate", help="JSON object, JSON file, or enumerate:<budget-index>:<result-index>")
    parser.add_argument("--profile-target", choices=("baseline", "candidate"), default="candidate")
    parser.add_argument("--workspace-bytes", type=int, default=32 << 20)
    parser.add_argument("--warmup", type=int, default=2)
    parser.add_argument("--rounds", type=int, default=6)
    parser.add_argument("--checkpoint-sha256", action="store_true", help="Expensive: actually calculate and record the checkpoint SHA-256")
    return parser.parse_args()


def load_extension(args: argparse.Namespace):
    from torch.utils.cpp_extension import CUDA_HOME, load

    if CUDA_HOME is None:
        raise RuntimeError("CUDA_HOME is unavailable; CUDA 12.9 or newer headers are required")
    kwargs: dict[str, Any] = {}
    if args.build_directory is not None:
        args.build_directory.mkdir(parents=True, exist_ok=True)
        kwargs["build_directory"] = str(args.build_directory)
    return load(
        name="h3_fc2_nvfp4_lt_schedule",
        sources=[str(SOURCE)],
        extra_include_paths=[str(Path(CUDA_HOME) / "include")],
        extra_cflags=["-O2", "-std=c++17"],
        extra_ldflags=["-L" + str(Path(CUDA_HOME) / "lib64"), "-lcublasLt", "-lcublas", "-lcudart"],
        verbose=args.verbose_build,
        **kwargs,
    )


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(8 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def tensor_sha256(tensor: torch.Tensor) -> str:
    immutable = tensor.detach().contiguous().clone().view(torch.uint16).cpu()
    return hashlib.sha256(immutable.numpy().tobytes()).hexdigest()


def compare(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, Any]:
    actual_copy = actual.detach().clone()
    expected_copy = expected.detach().clone()
    different = int(torch.count_nonzero(actual_copy != expected_copy))
    delta = (actual_copy.float() - expected_copy.float()).abs()
    return {
        "bf16_exact": different == 0,
        "different_elements": different,
        "max_abs": float(delta.max()),
        "mean_abs": float(delta.mean()),
        "actual_sha256": tensor_sha256(actual_copy),
        "expected_sha256": tensor_sha256(expected_copy),
    }


def percentile(values: list[float], fraction: float) -> float:
    ordered = sorted(values)
    if len(ordered) == 1:
        return ordered[0]
    position = (len(ordered) - 1) * fraction
    lower = int(position)
    weight = position - lower
    return ordered[lower] * (1.0 - weight) + ordered[min(lower + 1, len(ordered) - 1)] * weight


def timing_summary(milliseconds: list[float], m: int | None = None, n: int | None = None, k: int | None = None) -> dict[str, Any]:
    result: dict[str, Any] = {
        "samples_ms": milliseconds,
        "p50_ms": percentile(milliseconds, 0.50),
        "p95_ms": percentile(milliseconds, 0.95),
        "mean_ms": statistics.fmean(milliseconds),
    }
    if m is not None and n is not None and k is not None:
        result["dense_tflop_s_p50"] = 2.0 * m * n * k / (result["p50_ms"] * 1.0e9)
    return result


def timed_cuda(call: Callable[[], torch.Tensor]) -> tuple[float, torch.Tensor]:
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    output = call()
    end.record()
    end.synchronize()
    return float(start.elapsed_time(end)), output


def environment(args: argparse.Namespace, extension) -> dict[str, Any]:
    checkpoint = Path(args.model_path)
    try:
        commit = subprocess.run(
            ["git", "rev-parse", "HEAD"], cwd=ROOT, check=True,
            capture_output=True, text=True,
        ).stdout.strip()
    except (OSError, subprocess.CalledProcessError):
        commit = None
    result = {
        "platform": platform.platform(),
        "python": sys.version,
        "torch": torch.__version__,
        "torch_cuda": torch.version.cuda,
        "device": torch.cuda.get_device_name(),
        "device_capability": list(torch.cuda.get_device_capability()),
        "driver": torch.cuda.driver_version() if hasattr(torch.cuda, "driver_version") else None,
        "git_commit": commit,
        "checkpoint_path": str(checkpoint),
        "checkpoint_sha256": sha256_file(checkpoint) if args.checkpoint_sha256 else None,
        "checkpoint_hash_note": "calculated" if args.checkpoint_sha256 else "not calculated",
        "environment_switches": {key: value for key, value in sorted(os.environ.items()) if key.startswith(ENV_PREFIXES)},
        "extension": dict(extension.build_info()),
    }
    try:
        import comfy_kitchen
        result["comfy_kitchen"] = getattr(comfy_kitchen, "__version__", "0.2.31 package without __version__")
    except Exception as error:
        result["comfy_kitchen"] = f"import error: {error}"
    return result


def make_workload(args: argparse.Namespace, checkpoint: H3Checkpoint, model: H3PackedDenoiser):
    torch.manual_seed(args.seed)
    if args.capture is not None:
        path = args.capture / "input.pt" if args.capture.is_dir() else args.capture
        payload = torch.load(path, map_location="cuda", weights_only=False)
        hidden = payload["hidden"].to("cuda").contiguous()
        timesteps = payload["timesteps"].to("cuda")
        positions = payload["position_ids"].to("cuda")
        segments = payload["segments"]
        metadata = {"capture": str(path)}
    else:
        packer = H3PromptPacker(checkpoint)
        video, audio, aligned_frames = random_av_latents(
            args.width, args.height, args.frames, args.seed, device="cuda",
        )
        sigma = beta_sigmas(args.steps, device="cuda")[args.sampler_step - 1]
        native_audio = audio.to(torch.bfloat16) * (_audio_sigma(sigma) / sigma)
        text = torch.randn(1, args.text_tokens, 5376, device="cuda", dtype=torch.bfloat16)
        hidden, timesteps, segments, positions, _, _ = packer(
            text, video, native_audio, _model_sigma(sigma),
        )
        metadata = {
            "resolution": [args.width, args.height],
            "frames": aligned_frames,
            "steps": args.steps,
            "sampler_step": args.sampler_step,
            "seed": args.seed,
            "text_tokens": args.text_tokens,
        }
        canonical = (args.width, args.height, args.frames, args.seed, args.sampler_step, args.text_tokens) == (1344, 768, 124, 440420, 1, 100)
        if canonical and hidden.shape[0] != 37_810:
            raise RuntimeError(f"canonical workload must contain 37,810 tokens, got {hidden.shape[0]}")
    rotation = h3_rope_rotation(positions.to("cuda"), model.backbone.inv_freq, torch.bfloat16)
    metadata.update({"tokens": hidden.shape[0], "hidden_shape": list(hidden.shape), "segments": segments})
    return hidden, timesteps, rotation, segments, metadata


def capture_boundaries(args: argparse.Namespace, model: H3PackedDenoiser, hidden, timesteps, rotation, segments):
    wanted = set(args.blocks)
    if wanted != set(DEFAULT_BLOCKS):
        missing = set(DEFAULT_BLOCKS) - wanted
        if missing:
            raise ValueError(f"--blocks must retain required blocks 0,24,49; missing {sorted(missing)}")
    block_inputs: dict[int, torch.Tensor] = {}
    block_outputs: dict[int, torch.Tensor] = {}
    gate_up: dict[int, torch.Tensor] = {}
    hooks = []
    modulated_forwards = {}
    for index in wanted:
        fc1 = model.backbone.blocks[index].mlp.fc1
        hooks.append(fc1.register_forward_hook(
            lambda _module, _inputs, output, index=index: gate_up.__setitem__(index, output.detach().clone())
        ))
        original = fc1.forward_modulated
        modulated_forwards[index] = original

        def capture_modulated(_self, *values, index=index, original=original, **kwargs):
            output = original(*values, **kwargs)
            gate_up[index] = output.detach().clone()
            return output

        fc1.forward_modulated = types.MethodType(capture_modulated, fc1)
    try:
        with torch.inference_mode():
            for index, (block, adaln) in enumerate(zip(model.backbone.blocks, model.backbone.adaln, strict=True)):
                if index in wanted:
                    block_inputs[index] = hidden.detach().clone()
                hidden = block(hidden, rotation, *adaln(timesteps), segments)
                if index in wanted:
                    block_outputs[index] = hidden.detach().clone()
    finally:
        for hook in hooks:
            hook.remove()
        for index, original in modulated_forwards.items():
            model.backbone.blocks[index].mlp.fc1.forward_modulated = original
    if set(gate_up) != wanted:
        raise RuntimeError(f"failed to capture FC1 boundaries: got {sorted(gate_up)}")
    return block_inputs, block_outputs, gate_up


def packed_boundary(gate_up: torch.Tensor, fc2):
    tensor_scale, qdata, block_scale = vortex_native_quantize_swiglu_nvfp4(gate_up)
    alpha = (tensor_scale.float() * fc2.weight_scale_2.float()).reshape(1).contiguous()
    beta = torch.zeros(1, device=gate_up.device, dtype=torch.float32)
    return tensor_scale, qdata, block_scale, alpha, beta


def baseline_metadata(fc2, gate_up: torch.Tensor) -> dict[str, Any]:
    from torch.profiler import ProfilerActivity, profile

    with torch.inference_mode():
        fc2.forward_swiglu(gate_up)
        torch.cuda.synchronize()
        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as captured:
            output = fc2.forward_swiglu(gate_up)
        torch.cuda.synchronize()
    events = []
    for event in captured.events():
        if "CUDA" not in str(getattr(event, "device_type", "")):
            continue
        device_us = max(
            float(getattr(event, "device_time_total", 0.0) or 0.0),
            float(getattr(event, "self_device_time_total", 0.0) or 0.0),
            float(getattr(event, "cuda_time_total", 0.0) or 0.0),
            float(getattr(event, "self_cuda_time_total", 0.0) or 0.0),
        )
        events.append({"name": event.name, "device_time_total_us": device_us})
    events.sort(key=lambda row: row["device_time_total_us"], reverse=True)
    return {
        "path": "fc2.forward_swiglu -> accepted producer -> Comfy Kitchen 0.2.31 scaled_mm_nvfp4",
        "descriptors": {
            "packed_input_output": "row-major [M,K] @ [N,K].T -> BF16 [M,N]",
            "block_scale_mode": "VEC16_UE4M3",
            "compute_and_scale": "FP32",
            "scalar_pointer_mode": "device",
            "bias": None,
            "beta": 0.0,
            "comfy_kitchen_version": "0.2.31",
        },
        "profiler_cuda_events_available": bool(events),
        "profiler_note": None if events else "Torch profiler returned no CUDA kernel events on this build; use --mode profile with NCU for kernel metadata.",
        "top_cuda_events": events[:10],
        "output_sha256": tensor_sha256(output.detach().clone()),
    }


def enumerate_all(extension, packed, fc2, args: argparse.Namespace):
    _tensor_scale, qdata, block_scale, _alpha, _beta = packed
    groups = []
    errors = []
    for budget in args.workspace_budgets:
        try:
            values = [dict(value) for value in extension.enumerate(
                qdata, block_scale, fc2.weight, fc2.weight_scale, budget, args.requested_count,
            )]
            groups.append({"max_workspace_bytes": budget, "requested_count": args.requested_count, "returned_count": len(values), "algorithms": values})
        except Exception as error:
            groups.append({"max_workspace_bytes": budget, "requested_count": args.requested_count, "returned_count": 0, "algorithms": []})
            errors.append({"operation": "heuristic", "max_workspace_bytes": budget, "error": repr(error)})
    return groups, errors


def explicit_checks(extension, packed, fc2, enumerated, args: argparse.Namespace):
    _tensor_scale, qdata, block_scale, _alpha, _beta = packed
    unique_bases: dict[int, dict[str, Any]] = {}
    for group in enumerated:
        for candidate in group["algorithms"]:
            unique_bases.setdefault(candidate["algorithm_id"], candidate)
    checked, errors = [], []
    attempts = 0
    for base in unique_bases.values():
        mask = int(base.get("capabilities", {}).get("reduction_scheme_mask", 0))
        reductions = [1 << bit for bit in range(32) if mask & (1 << bit)]
        for split_k in args.explicit_split_k:
            schemes = [0] if split_k == 1 else reductions
            if not schemes:
                errors.append({"algorithm_id": base["algorithm_id"], "split_k": split_k, "error": "capability reports no public reduction scheme"})
            for reduction in schemes:
                if attempts >= args.max_explicit_checks:
                    return checked, errors
                attempts += 1
                config = {key: base[key] for key in ("algorithm_id", "tile_id", "stages_id", "custom_option", "cta_swizzle", "inner_shape", "cluster_shape") if key in base}
                config.update({"split_k": split_k, "reduction_scheme": reduction})
                try:
                    result = dict(extension.check(qdata, block_scale, fc2.weight, fc2.weight_scale, config))
                    result["requested_config"] = config
                    checked.append(result)
                except Exception as error:
                    errors.append({"config": config, "error": repr(error)})
    return checked, errors


def resolve_candidate(args: argparse.Namespace, enumerated, explicit) -> dict[str, Any]:
    if args.candidate:
        if args.candidate.startswith("enumerate:"):
            _, group, result = args.candidate.split(":")
            return dict(enumerated[int(group)]["algorithms"][int(result)])
        path = Path(args.candidate)
        return json.loads(path.read_text()) if path.exists() else json.loads(args.candidate)
    valid_explicit = [row for row in explicit if row.get("valid") and row.get("required_workspace_bytes", sys.maxsize) <= args.workspace_bytes]
    if valid_explicit:
        return valid_explicit[0]["requested_config"]
    for group in enumerated:
        for row in group["algorithms"]:
            if row.get("valid") and row.get("required_workspace_bytes", sys.maxsize) <= args.workspace_bytes:
                return row
    raise RuntimeError("no valid candidate fits --workspace-bytes")


def candidate_runner(extension, packed, fc2, config, workspace_bytes: int):
    _tensor_scale, qdata, block_scale, alpha, beta = packed
    output = torch.empty((qdata.shape[0], fc2.out_features), device=qdata.device, dtype=torch.bfloat16)
    workspace = torch.empty(workspace_bytes, device=qdata.device, dtype=torch.uint8)

    def run() -> torch.Tensor:
        extension.run(qdata, block_scale, fc2.weight, fc2.weight_scale, alpha, beta, output, workspace, config)
        return output

    return run, output, workspace


def comfy_packed_runner(packed, fc2, logical_rows: int):
    import torch.nn.functional as functional
    from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout

    tensor_scale, qdata, block_scale, _alpha, _beta = packed
    activation = QuantizedTensor(
        qdata,
        "TensorCoreNVFP4Layout",
        TensorCoreNVFP4Layout.Params(
            scale=tensor_scale,
            orig_dtype=torch.bfloat16,
            orig_shape=(logical_rows, fc2.in_features),
            block_scale=block_scale,
        ),
    )
    weight = fc2._packed_weight()

    def run():
        return functional.linear(activation, weight, None)[:logical_rows, :fc2.out_features]

    return run


def benchmark_pair(baseline, candidate, rounds: int, warmup: int, m: int, n: int, k: int, *, report_dense: bool = True):
    with torch.inference_mode():
        for _ in range(warmup):
            baseline()
            candidate()
        torch.cuda.synchronize()
        samples = {"baseline": [], "candidate": []}
        last = {}
        for round_index in range(rounds):
            order = (("baseline", baseline), ("candidate", candidate)) if round_index % 2 == 0 else (("candidate", candidate), ("baseline", baseline))
            for name, call in order:
                elapsed, output = timed_cuda(call)
                samples[name].append(elapsed)
                last[name] = output.detach().clone()
    return {
        "order": "AB/BA alternates by round",
        "baseline": timing_summary(samples["baseline"], m if report_dense else None, n if report_dense else None, k if report_dense else None),
        "candidate": timing_summary(samples["candidate"], m if report_dense else None, n if report_dense else None, k if report_dense else None),
        "parity": compare(last["candidate"][:m, :n], last["baseline"][:m, :n]),
    }


def benchmark_fc2_candidate(extension, gate_up, fc2, packed, config, args):
    run_candidate, output, workspace = candidate_runner(extension, packed, fc2, config, args.workspace_bytes)
    m, n, k = gate_up.shape[0], fc2.out_features, fc2.in_features
    run_comfy_packed = comfy_packed_runner(packed, fc2, m)
    direct = benchmark_pair(
        run_comfy_packed, run_candidate,
        args.rounds, args.warmup, m, n, k,
    )

    boundary_output = torch.empty_like(output)
    boundary_workspace = torch.empty_like(workspace)

    def candidate_boundary():
        fresh = packed_boundary(gate_up, fc2)
        extension.run(
            fresh[1], fresh[2], fc2.weight, fc2.weight_scale, fresh[3], fresh[4],
            boundary_output, boundary_workspace, config,
        )
        return boundary_output

    boundary = benchmark_pair(
        lambda: fc2.forward_swiglu(gate_up), candidate_boundary,
        args.rounds, args.warmup, m, n, k,
    )
    if not direct["parity"]["bf16_exact"] or not boundary["parity"]["bf16_exact"]:
        raise RuntimeError("FC2 library candidate is not byte-exact")
    checked = dict(extension.check(packed[1], packed[2], fc2.weight, fc2.weight_scale, config))
    return {
        "selected_config": config,
        "supplied_workspace_bytes": workspace.numel(),
        "required_workspace_bytes": checked.get("required_workspace_bytes"),
        "checked": checked,
        "fc2_only": direct,
        "accepted_producer_plus_fc2": boundary,
        "output_sha256": tensor_sha256(output[:m].detach().clone()),
    }


def profile_one(extension, packed, fc2, config, logical_rows: int, args):
    run_candidate, output, workspace = candidate_runner(extension, packed, fc2, config, args.workspace_bytes)
    expected = None
    if args.profile_target == "baseline":
        run_profiled = comfy_packed_runner(packed, fc2, logical_rows)
    else:
        run_profiled = run_candidate
    with torch.inference_mode():
        for _ in range(args.warmup):
            run_profiled()
        if args.profile_target == "candidate":
            expected = comfy_packed_runner(packed, fc2, logical_rows)().detach().clone()
        torch.cuda.synchronize()
        torch.cuda.cudart().cudaProfilerStart()
        output = run_profiled()
        torch.cuda.cudart().cudaProfilerStop()
        torch.cuda.synchronize()
    parity = compare(output[:logical_rows], expected[:logical_rows]) if expected is not None else None
    if parity is not None and not parity["bf16_exact"]:
        raise RuntimeError("profiled FC2 library candidate is not byte-exact")
    return {
        "target": args.profile_target,
        "selected_config": config if args.profile_target == "candidate" else None,
        "supplied_workspace_bytes": workspace.numel() if args.profile_target == "candidate" else 32 << 20,
        "output_sha256": tensor_sha256(output[:logical_rows, :fc2.out_features]),
        "candidate_vs_baseline": parity,
    }


def block_gate(extension, model, block_inputs, block_outputs, gate_up, timesteps, rotation, segments, config, args):
    rows = []
    for index in DEFAULT_BLOCKS:
        block = model.backbone.blocks[index]
        fc2 = block.mlp.fc2
        packed = packed_boundary(gate_up[index], fc2)
        _tensor_scale, qdata, _block_scale, _alpha, beta = packed
        candidate_output = torch.empty(
            (qdata.shape[0], fc2.out_features), device=qdata.device, dtype=torch.bfloat16,
        )
        workspace = torch.empty(args.workspace_bytes, device=qdata.device, dtype=torch.uint8)
        original = fc2.forward_swiglu

        def replacement(_self, actual_gate_up, expected=gate_up[index]):
            if actual_gate_up.shape != expected.shape:
                raise RuntimeError("block-gate FC2 received an unexpected boundary shape")
            tensor_scale, actual_qdata, actual_block_scale = vortex_native_quantize_swiglu_nvfp4(actual_gate_up)
            alpha = (tensor_scale.float() * fc2.weight_scale_2.float()).reshape(1).contiguous()
            extension.run(
                actual_qdata, actual_block_scale, fc2.weight, fc2.weight_scale,
                alpha, beta, candidate_output, workspace, config,
            )
            return candidate_output[:actual_gate_up.shape[0], :fc2.out_features]

        candidate_method = types.MethodType(replacement, fc2)
        adaln_values = tuple(value.detach().clone() for value in model.backbone.adaln[index](timesteps))

        def baseline():
            fc2.forward_swiglu = original
            return block(block_inputs[index].detach().clone(), rotation, *adaln_values, segments)

        def candidate():
            fc2.forward_swiglu = candidate_method
            return block(block_inputs[index].detach().clone(), rotation, *adaln_values, segments)

        try:
            timing = benchmark_pair(
                baseline, candidate, args.rounds, args.warmup,
                block_outputs[index].shape[0], block_outputs[index].shape[1], 1,
                report_dense=False,
            )
            baseline_value = baseline().detach().clone()
            candidate_value = candidate().detach().clone()
        finally:
            fc2.forward_swiglu = original
        row = {
            "block": index,
            "only_monkeypatched_method": "block.mlp.fc2.forward_swiglu",
            "accepted_gate_and_residual_path_preserved": True,
            "candidate_vs_baseline": compare(candidate_value, baseline_value),
            "baseline_vs_traversal": compare(baseline_value, block_outputs[index]),
            "candidate_vs_traversal": compare(candidate_value, block_outputs[index]),
            "timing": timing,
            "supplied_workspace_bytes": workspace.numel(),
        }
        if not all(
            row[name]["bf16_exact"]
            for name in ("candidate_vs_baseline", "baseline_vs_traversal", "candidate_vs_traversal")
        ):
            raise RuntimeError(f"FC2 library candidate is not byte-exact in block {index}")
        rows.append(row)
    return rows


def trajectory_gate(extension, checkpoint, model, config, args):
    packer = H3PromptPacker(checkpoint)
    torch.manual_seed(args.seed)
    video, audio, aligned_frames = random_av_latents(
        args.width, args.height, args.frames, args.seed, device="cuda",
    )
    text = torch.randn(
        1, args.text_tokens, 5376, device="cuda", dtype=torch.bfloat16,
    )
    originals = [block.mlp.fc2.forward_swiglu for block in model.backbone.blocks]
    shared_output = torch.empty(
        (((37_810 + 15) // 16) * 16, 5376), device="cuda", dtype=torch.bfloat16,
    )
    workspace = torch.empty(args.workspace_bytes, device="cuda", dtype=torch.uint8)
    beta = torch.zeros(1, device="cuda", dtype=torch.float32)

    def candidate_method(fc2):
        def replacement(_self, actual_gate_up):
            tensor_scale, qdata, block_scale = vortex_native_quantize_swiglu_nvfp4(actual_gate_up)
            if qdata.shape[0] > shared_output.shape[0] or fc2.out_features > shared_output.shape[1]:
                raise RuntimeError("trajectory FC2 boundary exceeds the preallocated canonical output")
            alpha = (tensor_scale.float() * fc2.weight_scale_2.float()).reshape(1).contiguous()
            output = shared_output[:qdata.shape[0], :fc2.out_features]
            extension.run(
                qdata, block_scale, fc2.weight, fc2.weight_scale,
                alpha, beta, output, workspace, config,
            )
            return output[:actual_gate_up.shape[0], :fc2.out_features]
        return types.MethodType(replacement, fc2)

    candidates = [candidate_method(block.mlp.fc2) for block in model.backbone.blocks]

    def run(candidate: bool):
        for index, block in enumerate(model.backbone.blocks):
            block.mlp.fc2.forward_swiglu = candidates[index] if candidate else originals[index]
        torch.cuda.synchronize()
        started = time.perf_counter()
        result = sample_video_res_multistep(
            model,
            packer,
            text,
            video.detach().clone(),
            audio.detach().clone(),
            steps=args.steps,
            seed=args.seed,
            return_audio=True,
            progress=True,
        )
        torch.cuda.synchronize()
        return result, time.perf_counter() - started

    try:
        with torch.inference_mode():
            (reference_video, reference_audio), baseline_seconds = run(False)
            (candidate_video, candidate_audio), candidate_seconds = run(True)
    finally:
        for block, original in zip(model.backbone.blocks, originals, strict=True):
            block.mlp.fc2.forward_swiglu = original

    video_parity = compare(candidate_video, reference_video)
    audio_parity = compare(candidate_audio, reference_audio)
    result = {
        "steps": args.steps,
        "resolution": [args.width, args.height],
        "frames": aligned_frames,
        "seed": args.seed,
        "baseline_seconds": baseline_seconds,
        "candidate_seconds": candidate_seconds,
        "improvement_percent": (1.0 - candidate_seconds / baseline_seconds) * 100.0,
        "video_parity": video_parity,
        "audio_parity": audio_parity,
        "bf16_exact": video_parity["bf16_exact"] and audio_parity["bf16_exact"],
        "selected_config": config,
        "supplied_workspace_bytes": workspace.numel(),
        "all_50_fc2_calls_replaced": True,
        "accepted_swiglu_producer_preserved": True,
        "accepted_gate_and_residual_path_preserved": True,
    }
    if not result["bf16_exact"]:
        raise RuntimeError("FC2 library candidate trajectory is not byte-exact")
    return result


def main() -> None:
    args = parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required")
    extension = load_extension(args)
    if args.mode == "compile":
        result = {"mode": "compile", "extension": dict(extension.build_info())}
        print(json.dumps(result, indent=2), flush=True)
        return

    checkpoint = H3Checkpoint(args.model_path, device="cuda")
    model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend="sage2").eval()
    hidden, timesteps, rotation, segments, workload = make_workload(args, checkpoint, model)
    with torch.inference_mode():
        block_inputs, block_outputs, gate_ups = capture_boundaries(
            args, model, hidden, timesteps, rotation, segments,
        )
    probe = args.probe_block
    if probe not in gate_ups:
        raise ValueError("--probe-block must be one of the retained blocks")
    fc2 = model.backbone.blocks[probe].mlp.fc2
    if fc2.bias is not None:
        raise RuntimeError("this no-bias FC2 scheduling study refuses a biased module")
    packed = packed_boundary(gate_ups[probe], fc2)
    enumerated, errors = enumerate_all(extension, packed, fc2, args)
    explicit, explicit_errors = explicit_checks(extension, packed, fc2, enumerated, args)
    errors.extend(explicit_errors)
    result: dict[str, Any] = {
        "mode": args.mode,
        "environment": environment(args, extension),
        "workload": workload,
        "retained_blocks": list(DEFAULT_BLOCKS),
        "immutable_cloned_block_inputs": {str(index): list(value.shape) for index, value in block_inputs.items()},
        "fc2_boundary": {
            "block": probe,
            "gate_up_shape": list(gate_ups[probe].shape),
            "activation_qdata_shape": list(packed[1].shape),
            "weight_qdata_shape": list(fc2.weight.shape),
            "logical_mnk": [gate_ups[probe].shape[0], fc2.out_features, fc2.in_features],
            "descriptor_mnk_after_padding": [packed[1].shape[0], fc2.out_features, fc2.in_features],
            "producer": "vortex_native_quantize_swiglu_nvfp4",
            "no_bias": fc2.bias is None,
        },
        "baseline_kernel_metadata": baseline_metadata(fc2, gate_ups[probe]),
        "heuristics": enumerated,
        "explicit_split_k_checks": explicit,
    }

    if args.mode in {"characterize", "sweep", "selected", "profile", "block-gate", "trajectory"}:
        config = resolve_candidate(args, enumerated, explicit)
        result["selected"] = config
    if args.mode == "characterize":
        result["candidate_probe"] = benchmark_fc2_candidate(extension, gate_ups[probe], fc2, packed, config, args)
    elif args.mode == "sweep":
        candidates = []
        seen = set()
        pool = [row for group in enumerated for row in group["algorithms"]] + [row for row in explicit if row.get("valid")]
        for row in pool:
            candidate = row.get("requested_config", row)
            key = tuple(candidate.get(name) for name in ("algorithm_id", "tile_id", "stages_id", "split_k", "reduction_scheme", "custom_option", "cta_swizzle", "inner_shape", "cluster_shape"))
            if key in seen or row.get("required_workspace_bytes", 0) > args.workspace_bytes:
                continue
            seen.add(key)
            try:
                candidates.append(benchmark_fc2_candidate(extension, gate_ups[probe], fc2, packed, candidate, args))
            except Exception as error:
                errors.append({"config": candidate, "operation": "benchmark", "error": repr(error)})
            if len(candidates) >= args.candidate_limit:
                break
        result["candidates"] = candidates
    elif args.mode == "selected":
        result["candidate_probe"] = benchmark_fc2_candidate(extension, gate_ups[probe], fc2, packed, config, args)
    elif args.mode == "profile":
        result["profile"] = profile_one(extension, packed, fc2, config, gate_ups[probe].shape[0], args)
    elif args.mode == "block-gate":
        result["block_gate"] = block_gate(extension, model, block_inputs, block_outputs, gate_ups, timesteps, rotation, segments, config, args)
    elif args.mode == "trajectory":
        result["trajectory"] = trajectory_gate(extension, checkpoint, model, config, args)

    result["errors_and_unsupported"] = errors + [{
        "feature": "Stream-K",
        "supported_public_control": False,
        "reason": extension.build_info()["stream_k_note"],
    }]
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result, indent=2), flush=True)


if __name__ == "__main__":
    main()