#!/usr/bin/env python3
"""Freeze vLLM 0.20 microprofiles into Frontier-compatible CSV inputs."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import re
import shutil
from pathlib import Path
from typing import Any


STAT_NAMES = ("min", "max", "mean", "median", "std")
ATTENTION_OPS = (
    "attn_input_reshape",
    "attn_kv_cache_save",
    "attn_prefill",
    "attn_decode",
    "attn_output_reshape",
)
ATTENTION_METADATA = (
    "n_embd",
    "n_q_head",
    "n_kv_head",
    "block_size",
    "num_tensor_parallel_workers",
    "max_model_len",
    "batch_size",
    "prefill_chunk_size",
    "kv_cache_size",
    "is_prefill",
    "attention_backend",
    "is_mixed_batch",
    "mode",
    "seq_lens",
    "total_tokens",
    "max_seq_len",
    "min_seq_len",
    "avg_seq_len",
    "equal_seq_len",
    "seq_len_variance",
    "seq_len_std",
    "seq_len_cv",
    "is_chunked_prefill_sample",
    "chunk_start_token",
    "chunk_end_token",
    "total_prefill_tokens",
    "profiling_precision",
    "model_arch",
    "quant_signature",
    "measurement_type",
    "is_true_mixed_batch",
    "prefill_seq_lens",
    "prefill_kv_cache_sizes",
    "decode_kv_cache_sizes",
    "num_prefill_seqs",
    "num_decode_seqs",
    "decode_batch_size",
    "total_batch_size",
    "total_decode_tokens",
    "decode_avg_kv_cache_size",
    "batch_composition_ratio",
    "batch_spec",
    "projection_policy",
)
MOE_OPS = (
    "moe_gating_linear",
    "moe_gating_routing_topk",
    "moe_shuffling",
    "moe_grouped_gemm",
)
MOE_METADATA = (
    "num_tokens",
    "num_experts",
    "num_experts_per_device",
    "expert_parallel_size",
    "routing_runtime_path",
    "routing_assignment_policy",
    "routing_weight_policy",
    "routing_uses_router_logits",
    "gating_runtime_context",
    "gating_runtime_context_impl",
    "router_topk",
    "hidden_dim",
    "expert_hidden_dim",
    "use_gated",
    "num_tensor_parallel_workers",
    "total_routed_tokens",
    "model_expansion_ratio",
    "tokens_per_expert_avg",
    "tokens_to_experts_ratio",
    "expert_utilization",
    "min_load_ratio",
    "load_imbalance_cv",
    "max_load_ratio",
    "load_entropy",
    "load_gini_coefficient",
    "load_distribution",
    "seed",
    "moe_grouped_gemm_backend",
    "measurement_type",
    "profiling_precision",
    "model_arch",
    "quant_signature",
    "router_median_nonadditivity_ratio",
    "projection_policy",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--linear", type=Path, required=True)
    parser.add_argument("--attention", type=Path, nargs="+", required=True)
    parser.add_argument("--moe", type=Path, required=True)
    parser.add_argument("--router", type=Path, required=True)
    parser.add_argument("--allreduce", type=Path, nargs=2)
    parser.add_argument("--allreduce-frozen", type=Path)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--measurement-type",
        choices=("CUDA_EVENT", "KERNEL_ONLY"),
        default="CUDA_EVENT",
    )
    parser.add_argument(
        "--frontier-commit",
        default="d9cfeb6d8791fbf2f295dd9744c56a666171776e",
    )
    return parser.parse_args()


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


def load_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text())


def stat_columns(prefix: str, stats: dict[str, float]) -> dict[str, float]:
    return {f"time_stats.{prefix}.{name}": float(stats[name]) for name in STAT_NAMES}


def zero_stat_columns(prefix: str) -> dict[str, float]:
    return {f"time_stats.{prefix}.{name}": 0.0 for name in STAT_NAMES}


def attention_core_stats(raw: dict[str, Any]) -> dict[str, float]:
    # vLLM's benchmark result exports aggregate mean but not the raw samples.
    # Preserve that mean as Frontier's training target and record the proxy in
    # the manifest rather than inventing an unobserved median.
    return {
        "min": 1000.0 * float(raw["min_time"]),
        "max": 1000.0 * float(raw["max_time"]),
        "mean": 1000.0 * float(raw["mean_time"]),
        "median": 1000.0 * float(raw["mean_time"]),
        "std": 1000.0 * float(raw["std_time"]),
    }


def kv_update_stats(raw: dict[str, Any]) -> dict[str, float]:
    stats = raw["kv_cache_update_time"]
    return {name: float(stats[f"{name}_ms"]) for name in STAT_NAMES}


def parse_size(value: str, suffix: str) -> int:
    return int(value) * (1024 if suffix == "k" else 1)


def parse_batch_spec(spec: str) -> list[tuple[int, int]]:
    requests: list[tuple[int, int]] = []
    pattern = re.compile(r"^(?:(\d+))?q(\d+)(k?)(?:s(\d+)(k?))?$")
    for segment in spec.split("_"):
        match = pattern.match(segment)
        if match is None:
            raise ValueError(f"invalid vLLM batch spec: {spec}")
        count = int(match.group(1) or 1)
        query = parse_size(match.group(2), match.group(3))
        kv = (
            parse_size(match.group(4), match.group(5))
            if match.group(4)
            else query
        )
        requests.extend([(query, kv)] * count)
    return requests


def write_csv(path: Path, fieldnames: list[str], rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as handle:
        writer = csv.DictWriter(
            handle,
            fieldnames=fieldnames,
            extrasaction="raise",
            lineterminator="\n",
        )
        writer.writeheader()
        writer.writerows(rows)


def freeze_attention(
    inputs: list[Path], output: Path, *, measurement_type: str
) -> tuple[int, int, list[str]]:
    rows: list[dict[str, Any]] = []
    mixed_rows: list[dict[str, Any]] = []
    seen_tps: set[int] = set()
    raw_by_tp: dict[int, list[dict[str, Any]]] = {}
    for path in inputs:
        payload = load_json(path)
        if payload.get("schema_version") != "qwen30_vllm020_flashattn_raw.v1":
            raise ValueError(f"unexpected attention schema in {path}")
        if payload["environment"].get("vllm_version") != "0.20.0":
            raise ValueError(f"unexpected vLLM version in {path}")
        expected_method = (
            "record_function" if measurement_type == "KERNEL_ONLY" else "cuda_event"
        )
        if payload["environment"].get("profile_method", "cuda_event") != expected_method:
            raise ValueError(
                f"attention profile method mismatch in {path}: expected {expected_method}"
            )
        for raw in payload["rows"]:
            if raw.get("error") is not None:
                raise ValueError(f"failed attention row in {path}: {raw['error']}")
            tp = int(raw["tensor_parallel_size"])
            seen_tps.add(tp)
            raw_by_tp.setdefault(tp, []).append(raw)

    def pure_reference_ms(
        tp: int, requests: list[tuple[int, int]], *, decode_phase: bool
    ) -> float:
        candidates: list[tuple[list[tuple[int, int]], float]] = []
        for candidate in raw_by_tp[tp]:
            parsed = parse_batch_spec(candidate["config"]["batch_spec"])
            is_decode = all(query == 1 for query, _ in parsed)
            if is_decode != decode_phase:
                continue
            if decode_phase and not is_decode:
                continue
            if not decode_phase and any(query == 1 for query, _ in parsed):
                continue
            candidates.append((parsed, 1000.0 * float(candidate["mean_time"])))
        for parsed, mean_ms in candidates:
            if parsed == requests:
                return mean_ms
        if not decode_phase:
            raise ValueError(f"no exact pure prefill reference for TP{tp}: {requests}")
        if len({kv for _, kv in requests}) != 1:
            raise ValueError(f"decode interpolation requires one KV length: {requests}")
        target_batch = len(requests)
        target_kv = requests[0][1]
        points = sorted(
            (parsed[0][1], mean_ms)
            for parsed, mean_ms in candidates
            if len(parsed) == target_batch
            and len({kv for _, kv in parsed}) == 1
        )
        if not points:
            raise ValueError(f"no pure decode reference for TP{tp}: {requests}")
        if target_kv <= points[0][0]:
            return points[0][1]
        if target_kv >= points[-1][0]:
            return points[-1][1]
        for (left_kv, left_ms), (right_kv, right_ms) in zip(points, points[1:]):
            if left_kv <= target_kv <= right_kv:
                fraction = (target_kv - left_kv) / (right_kv - left_kv)
                return left_ms + fraction * (right_ms - left_ms)
        raise AssertionError("unreachable decode interpolation")

    for tp in sorted(raw_by_tp):
        for raw in raw_by_tp[tp]:
            spec = raw["config"]["batch_spec"]
            requests = parse_batch_spec(spec)
            prefill = [(q, kv) for q, kv in requests if q > 1]
            decode = [(q, kv) for q, kv in requests if q == 1]
            core = attention_core_stats(raw)
            kv_stats = kv_update_stats(raw)
            if prefill and decode:
                prefill_reference_ms = pure_reference_ms(
                    tp, prefill, decode_phase=False
                )
                decode_reference_ms = pure_reference_ms(
                    tp, decode, decode_phase=True
                )
                reference_total_ms = prefill_reference_ms + decode_reference_ms
                prefill_share = prefill_reference_ms / reference_total_ms
                decode_share = decode_reference_ms / reference_total_ms
                projected_prefill = {
                    name: value * prefill_share for name, value in core.items()
                }
                projected_decode = {
                    name: value * decode_share for name, value in core.items()
                }
                row = {}
                for op in ATTENTION_OPS:
                    row.update(zero_stat_columns(op))
                row.update(stat_columns("attn_kv_cache_save", kv_stats))
                row.update(stat_columns("attn_prefill", projected_prefill))
                row.update(stat_columns("attn_decode", projected_decode))
                prefill_queries = [q for q, _ in prefill]
                prefill_contexts = [kv - q for q, kv in prefill]
                decode_kv_lengths = [kv for _, kv in decode]
                total_batch = len(requests)
                row.update(
                    {
                        "n_embd": 2048,
                        "n_q_head": 32,
                        "n_kv_head": 4,
                        "block_size": 16,
                        "num_tensor_parallel_workers": tp,
                        "max_model_len": 40960,
                        "batch_size": total_batch,
                        "prefill_chunk_size": 0,
                        "kv_cache_size": 0,
                        "is_prefill": True,
                        "attention_backend": "FLASH_ATTN",
                        "is_mixed_batch": False,
                        "mode": "true_mixed_fused_projected",
                        "seq_lens": "",
                        "total_tokens": sum(prefill_queries) + len(decode),
                        "max_seq_len": "",
                        "min_seq_len": "",
                        "avg_seq_len": "",
                        "equal_seq_len": "",
                        "seq_len_variance": "",
                        "seq_len_std": "",
                        "seq_len_cv": "",
                        "is_chunked_prefill_sample": False,
                        "chunk_start_token": "",
                        "chunk_end_token": "",
                        "total_prefill_tokens": sum(prefill_queries),
                        "profiling_precision": "BF16",
                        "model_arch": "generic",
                        "quant_signature": "none",
                        "measurement_type": measurement_type,
                        "is_true_mixed_batch": True,
                        "prefill_seq_lens": json.dumps(prefill_queries),
                        "prefill_kv_cache_sizes": json.dumps(prefill_contexts),
                        "decode_kv_cache_sizes": json.dumps(decode_kv_lengths),
                        "num_prefill_seqs": len(prefill),
                        "num_decode_seqs": len(decode),
                        "decode_batch_size": len(decode),
                        "total_batch_size": total_batch,
                        "total_decode_tokens": len(decode),
                        "decode_avg_kv_cache_size": (
                            sum(decode_kv_lengths) / len(decode_kv_lengths)
                        ),
                        "batch_composition_ratio": len(prefill) / total_batch,
                        "batch_spec": spec,
                        "projection_policy": (
                            "fused_total_conserving_projection_by_same_tp_pure_"
                            "prefill_decode_reference_ratio"
                        ),
                    }
                )
                rows.append(row)
                mixed_rows.append(
                    {
                        "num_tensor_parallel_workers": tp,
                        "batch_spec": spec,
                        "num_prefill_seqs": len(prefill),
                        "num_decode_seqs": len(decode),
                        "total_prefill_tokens": sum(q for q, _ in prefill),
                        "total_decode_tokens": len(decode),
                        "decode_avg_kv_cache_size": sum(kv for _, kv in decode)
                        / len(decode),
                        "attention_core_mean_ms": core["mean"],
                        "attention_core_mean_as_median_ms": core["median"],
                        "kv_cache_update_median_ms": kv_stats["median"],
                        "pure_prefill_reference_mean_ms": prefill_reference_ms,
                        "pure_decode_reference_mean_ms": decode_reference_ms,
                        "projected_prefill_mean_ms": projected_prefill["mean"],
                        "projected_decode_mean_ms": projected_decode["mean"],
                        "projection_sum_error_ms": (
                            projected_prefill["mean"]
                            + projected_decode["mean"]
                            - core["mean"]
                        ),
                        "representation": (
                            "one_fused_FA3_call_projected_for_Frontier_with_"
                            "total_conservation"
                        ),
                    }
                )
                continue

            is_decode = bool(decode)
            queries = [q for q, _ in requests]
            contexts = [kv if is_decode else kv - q for q, kv in requests]
            avg_query = sum(queries) / len(queries)
            variance = sum((query - avg_query) ** 2 for query in queries) / len(queries)
            std = math.sqrt(variance)
            avg_context = sum(contexts) / len(contexts)
            row: dict[str, Any] = {}
            for op in ATTENTION_OPS:
                row.update(zero_stat_columns(op))
            row.update(stat_columns("attn_kv_cache_save", kv_stats))
            row.update(
                stat_columns("attn_decode" if is_decode else "attn_prefill", core)
            )
            row.update(
                {
                    "n_embd": 2048,
                    "n_q_head": 32,
                    "n_kv_head": 4,
                    "block_size": 16,
                    "num_tensor_parallel_workers": tp,
                    "max_model_len": 40960,
                    "batch_size": len(requests),
                    "prefill_chunk_size": 0 if is_decode else sum(queries),
                    "kv_cache_size": avg_context,
                    "is_prefill": not is_decode,
                    "attention_backend": "FLASH_ATTN",
                    "is_mixed_batch": False,
                    "mode": "vllm020_batch_spec",
                    "seq_lens": json.dumps(queries),
                    "total_tokens": sum(queries),
                    "max_seq_len": max(queries),
                    "min_seq_len": min(queries),
                    "avg_seq_len": avg_query,
                    "equal_seq_len": len(set(queries)) == 1,
                    "seq_len_variance": variance,
                    "seq_len_std": std,
                    "seq_len_cv": std / avg_query if avg_query else 0.0,
                    "is_chunked_prefill_sample": (not is_decode and avg_context > 0),
                    "chunk_start_token": avg_context if not is_decode else 0,
                    "chunk_end_token": avg_context + sum(queries) if not is_decode else 0,
                    "total_prefill_tokens": 0 if is_decode else sum(queries),
                    "profiling_precision": "BF16",
                    "model_arch": "generic",
                    "quant_signature": "none",
                    "measurement_type": measurement_type,
                    "is_true_mixed_batch": False,
                    "prefill_seq_lens": "",
                    "prefill_kv_cache_sizes": "",
                    "decode_kv_cache_sizes": "",
                    "num_prefill_seqs": "",
                    "num_decode_seqs": "",
                    "decode_batch_size": "",
                    "total_batch_size": "",
                    "total_decode_tokens": "",
                    "decode_avg_kv_cache_size": "",
                    "batch_composition_ratio": "",
                    "batch_spec": spec,
                    "projection_policy": (
                        "measured_FA3_core_plus_measured_KV;reshape_assumed_zero;"
                        "mean_as_median"
                    ),
                }
            )
            rows.append(row)

    if seen_tps != {1, 2, 4}:
        raise ValueError(f"attention TP coverage mismatch: {seen_tps}")
    attention_fields = [
        f"time_stats.{op}.{stat}" for op in ATTENTION_OPS for stat in STAT_NAMES
    ] + list(ATTENTION_METADATA)
    write_csv(output / "attention.csv", attention_fields, rows)
    mixed_fields = [
        "num_tensor_parallel_workers",
        "batch_spec",
        "num_prefill_seqs",
        "num_decode_seqs",
        "total_prefill_tokens",
        "total_decode_tokens",
        "decode_avg_kv_cache_size",
        "attention_core_mean_ms",
        "attention_core_mean_as_median_ms",
        "kv_cache_update_median_ms",
        "pure_prefill_reference_mean_ms",
        "pure_decode_reference_mean_ms",
        "projected_prefill_mean_ms",
        "projected_decode_mean_ms",
        "projection_sum_error_ms",
        "representation",
    ]
    write_csv(output / "attention_true_mixed_fused.csv", mixed_fields, mixed_rows)
    return len(rows), len(mixed_rows), sorted(seen_tps)


def load_features(counts: list[int]) -> dict[str, float]:
    total = sum(counts)
    count = len(counts)
    mean = total / count
    variance = sum((value - mean) ** 2 for value in counts) / count
    probabilities = [value / total for value in counts if value > 0]
    entropy = -sum(probability * math.log2(probability) for probability in probabilities)
    sorted_counts = sorted(counts)
    gini = (
        2 * sum((index + 1) * value for index, value in enumerate(sorted_counts))
        / (count * total)
        - (count + 1) / count
    )
    return {
        "total_routed_tokens": total,
        "num_experts_per_device": count,
        "hidden_dim": 2048,
        "expert_hidden_dim": 768,
        "router_topk": 8,
        "model_expansion_ratio": 768 / 2048,
        "tokens_per_expert_avg": mean,
        "tokens_to_experts_ratio": mean,
        "expert_utilization": sum(value > 0 for value in counts) / count,
        "min_load_ratio": min(counts) / mean,
        "load_imbalance_cv": math.sqrt(variance) / mean,
        "max_load_ratio": max(counts) / mean,
        "load_entropy": entropy,
        "load_gini_coefficient": gini,
    }


def freeze_moe(
    moe_path: Path, router_path: Path, output: Path, *, measurement_type: str
) -> int:
    moe = load_json(moe_path)
    router = load_json(router_path)
    if moe.get("schema_version") != "qwen30_vllm020_moe_raw.v1":
        raise ValueError(f"unexpected MoE schema in {moe_path}")
    if router.get("schema_version") != "qwen30_vllm020_router_raw.v1":
        raise ValueError(f"unexpected router schema in {router_path}")
    expected_method = (
        "record_function" if measurement_type == "KERNEL_ONLY" else "cuda_event"
    )
    for payload, label in ((moe, "moe"), (router, "router")):
        if payload["environment"].get("profile_method", "cuda_event") != expected_method:
            raise ValueError(
                f"{label} profile method mismatch in {payload}: expected {expected_method}"
            )
    router_by_tokens = {int(row["num_tokens"]): row for row in router["rows"]}
    rows: list[dict[str, Any]] = []
    seen_pairs: set[tuple[int, int, str]] = set()
    for raw in moe["rows"]:
        tp = int(raw["tensor_parallel_size"])
        num_tokens = int(raw["num_tokens"])
        routing_mode = str(raw["routing_mode"])
        key = (tp, num_tokens, routing_mode)
        if key in seen_pairs:
            raise ValueError(f"duplicate MoE row: {key}")
        seen_pairs.add(key)
        router_row = router_by_tokens[num_tokens]
        counts = [int(value) for value in raw["routing_load"]["counts"]]
        if sum(counts) != num_tokens * 8 or len(counts) != 128:
            raise ValueError(f"invalid routing counts for {key}")
        row: dict[str, Any] = {}
        row.update(stat_columns("moe_gating_linear", router_row["gate_linear_time_ms"]))
        row.update(
            stat_columns(
                "moe_gating_routing_topk", router_row["routing_topk_time_ms"]
            )
        )
        row.update(zero_stat_columns("moe_shuffling"))
        row.update(stat_columns("moe_grouped_gemm", raw["time_ms"]))
        row.update(load_features(counts))
        row.update(
            {
                "num_tokens": num_tokens,
                "num_experts": 128,
                "expert_parallel_size": 1,
                "routing_runtime_path": "standard_fused_topk",
                "routing_assignment_policy": (
                    "logit_topk"
                    if routing_mode == "uniform_random_logits"
                    else "fixed_hotset8"
                ),
                "routing_weight_policy": "softmax_renorm",
                "routing_uses_router_logits": routing_mode == "uniform_random_logits",
                "gating_runtime_context": "standalone_legacy",
                "gating_runtime_context_impl": "vllm020_replicated_linear",
                "use_gated": True,
                "num_tensor_parallel_workers": tp,
                "load_distribution": routing_mode,
                "seed": 20260716,
                "moe_grouped_gemm_backend": raw["backend"],
                "measurement_type": measurement_type,
                "profiling_precision": "BF16",
                "model_arch": "generic",
                "quant_signature": "none",
                "router_median_nonadditivity_ratio": router_row[
                    "median_nonadditivity_ratio"
                ],
                "projection_policy": (
                    "measured_gate+topk+modular_expert;shuffling_zero_because_"
                    "expert_measurement_includes_prepare_finalize"
                ),
            }
        )
        rows.append(row)

    tokens = {int(row["num_tokens"]) for row in router["rows"]}
    modes = {str(row["routing_mode"]) for row in moe["rows"]}
    expected = {(tp, tokens_value, mode) for tp in (1, 2, 4) for tokens_value in tokens for mode in modes}
    actual = {
        (int(row["num_tensor_parallel_workers"]), int(row["num_tokens"]), str(row["load_distribution"]))
        for row in rows
    }
    if actual != expected:
        raise ValueError(
            f"MoE TP/token/routing coverage mismatch: missing={expected - actual}, extra={actual - expected}"
        )
    moe_fields = [
        f"time_stats.{op}.{stat}" for op in MOE_OPS for stat in STAT_NAMES
    ] + list(MOE_METADATA)
    write_csv(output / "moe.csv", moe_fields, rows)
    return len(rows)


def freeze_allreduce(inputs: list[Path], output: Path) -> int:
    rows: list[dict[str, Any]] = []
    environments: list[dict[str, Any]] = []
    for path in inputs:
        payload = load_json(path)
        if payload.get("schema_version") != "qwen30_vllm020_allreduce_raw.v1":
            raise ValueError(f"unexpected all-reduce schema in {path}")
        rows.extend(payload["rows"])
        environments.append(payload["environment"])
    if {(row["tensor_parallel_size"], row["num_tokens"]) for row in rows} != {
        (tp, tokens)
        for tp in (2, 4)
        for tokens in (1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192)
    }:
        raise ValueError("all-reduce TP/token coverage mismatch")
    (output / "allreduce.json").write_text(
        json.dumps(
            {
                "schema_version": "qwen30_vllm020_allreduce_frozen.v1",
                "environment": environments,
                "rows": sorted(
                    rows,
                    key=lambda row: (
                        row["tensor_parallel_size"],
                        row["num_tokens"],
                    ),
                ),
                "frontier_consumption": (
                    "diagnostic_only_in_base_profile_only_run; measured lookup "
                    "requires a separate CC-backend injection ablation"
                ),
            },
            indent=2,
            sort_keys=True,
        )
        + "\n"
    )
    return len(rows)


def main() -> None:
    args = parse_args()
    if args.allreduce is not None and args.allreduce_frozen is not None:
        raise SystemExit("provide either --allreduce or --allreduce-frozen, not both")
    all_inputs = [args.linear, *args.attention, args.moe, args.router]
    if args.allreduce is not None:
        all_inputs.extend(args.allreduce)
    if args.allreduce_frozen is not None:
        all_inputs.append(args.allreduce_frozen)
    for path in all_inputs:
        if not path.is_file():
            raise SystemExit(f"missing input: {path}")
    args.output.mkdir(parents=True, exist_ok=False)

    linear_output = args.output / "linear_op.csv"
    shutil.copyfile(args.linear, linear_output)
    with linear_output.open(newline="") as handle:
        linear_rows = list(csv.DictReader(handle))
    if not linear_rows:
        raise ValueError("linear profile has no rows")
    if {row.get("measurement_type") for row in linear_rows} != {args.measurement_type}:
        raise ValueError(
            f"linear measurement family mismatch: expected {args.measurement_type}"
        )

    attention_rows, mixed_rows, attention_tps = freeze_attention(
        list(args.attention), args.output, measurement_type=args.measurement_type
    )
    moe_rows = freeze_moe(
        args.moe, args.router, args.output, measurement_type=args.measurement_type
    )
    allreduce_rows = 0
    allreduce_source = "not_included"
    if args.allreduce is not None:
        allreduce_rows = freeze_allreduce(list(args.allreduce), args.output)
        allreduce_source = "raw_vllm020_measurements"
    elif args.allreduce_frozen is not None:
        shutil.copyfile(args.allreduce_frozen, args.output / "allreduce.json")
        allreduce_rows = len(load_json(args.allreduce_frozen).get("rows", []))
        allreduce_source = "carried_forward_frozen_measurements"

    output_files = [
        linear_output,
        args.output / "attention.csv",
        args.output / "attention_true_mixed_fused.csv",
        args.output / "moe.csv",
    ]
    if (args.output / "allreduce.json").is_file():
        output_files.append(args.output / "allreduce.json")
    batch_composition_augmented = len(args.attention) > 3
    long_context_augmented = args.measurement_type == "KERNEL_ONLY" or any(
        "long-context" in path.name for path in args.attention
    )
    long_context_coverage: dict[str, Any] = {"included": long_context_augmented}
    if long_context_augmented:
        with (args.output / "attention.csv").open(newline="") as handle:
            frozen_attention = list(csv.DictReader(handle))
        by_tp = {}
        for tp in (1, 2, 4):
            decode_kv = {
                int(float(row["kv_cache_size"]))
                for row in frozen_attention
                if int(row["num_tensor_parallel_workers"]) == tp
                and row["is_prefill"].lower() == "false"
            }
            required_decode = (
                {128, 1024, 2048, 4096, 8192, 16384, 32768, 40960}
                if args.measurement_type == "KERNEL_ONLY"
                else {16384, 32768, 40960}
            )
            if not required_decode.issubset(decode_kv):
                raise ValueError(f"decode KV coverage mismatch for TP{tp}")
            by_tp[str(tp)] = {
                "decode_kv_lengths": sorted(decode_kv),
            }
            if args.measurement_type != "KERNEL_ONLY":
                mixed_kv = {
                    int(float(row["decode_avg_kv_cache_size"]))
                    for row in frozen_attention
                    if int(row["num_tensor_parallel_workers"]) == tp
                    and row.get("is_true_mixed_batch", "").lower() == "true"
                }
                if not {16384, 32768}.issubset(mixed_kv):
                    raise ValueError(f"long-context mixed coverage mismatch for TP{tp}")
                by_tp[str(tp)]["true_mixed_decode_avg_kv_lengths"] = sorted(mixed_kv)
        long_context_coverage["by_tp"] = by_tp
    manifest = {
        "schema_version": (
            "frontier_qwen30_vllm020_kernel_only_profile.v1"
            if args.measurement_type == "KERNEL_ONLY"
            else (
                "frontier_qwen30_vllm020_frozen_profile.v4"
                if long_context_augmented
            else (
                "frontier_qwen30_vllm020_frozen_profile.v3"
                if batch_composition_augmented
                else "frontier_qwen30_vllm020_frozen_profile.v2"
                )
            )
        ),
        "profile_id": (
            "qwen3-30b-a3b-bf16-vllm020-h20-tp1-2-4-"
            "fused-mixed-total-conserving"
            + ("-kernel-only-record-function" if args.measurement_type == "KERNEL_ONLY" else "")
            + ("-pure-prefill-batch-composition" if batch_composition_augmented else "")
            + ("-long-context-decode-mixed" if long_context_augmented else "")
        ),
        "environment_contract": {
            "hardware": "NVIDIA H20",
            "model": "Qwen3-30B-A3B",
            "dtype": "bfloat16",
            "vllm_version": "0.20.0",
            "vllm_source_commit": "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1",
            "frontier_commit": args.frontier_commit,
            "tensor_parallel_sizes": [1, 2, 4],
        },
        "row_counts": {
            "linear": len(linear_rows),
            "attention_frontier_compatible": attention_rows,
            "attention_true_mixed_fused_diagnostic": mixed_rows,
            "moe": moe_rows,
            "allreduce": allreduce_rows,
        },
        "attention_tp_coverage": attention_tps,
        "attention_long_context_coverage": long_context_coverage,
        "projection_contract": {
            "linear": "Frontier profiler using vLLM 0.20 CUDA operators",
            "attention": (
                "Pure prefill/extend/decode FA3 core plus separately measured KV update; "
                "input/output reshape assumed zero; exported mean is used as median target; "
                "true mixed rows use a total-conserving compatibility projection"
            ),
            "attention_true_mixed": (
                "The directly measured fused total is preserved in diagnostics. Frontier's "
                "two targets are projected by the same-TP pure prefill/decode reference "
                "ratio, with projected prefill + decode exactly equal to the fused total; "
                "the split is a schema compatibility attribution, not an observation"
            ),
            "attention_pure_prefill_batch_composition": (
                "Direct FA3 measurements for 2/4 requests at query length 2048 and "
                "2/4/8/16 requests at query length 512 for each TP; included only "
                "when batch-composition attention inputs are supplied"
                if batch_composition_augmented
                else "not included"
            ),
            "attention_long_context": (
                "Direct FA3 measurements at decode KV lengths 16K/32K/40K, "
                "plus true-mixed prefill/decode measurements at 16K/32K, "
                "for TP1/2/4"
                if long_context_augmented
                else "not included"
            ),
            "moe": (
                "Replicated gate and fused top-k plus TP-local modular expert kernel; "
                "expert measurement already includes prepare/finalize so shuffling is zero"
            ),
            "allreduce": (
                "Frozen exact runtime measurements; source=" + allreduce_source
            ),
        },
        "inputs": {str(path.resolve()): sha256(path) for path in all_inputs},
        "outputs": {path.name: sha256(path) for path in output_files},
    }
    manifest_path = args.output / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    print(json.dumps(manifest["row_counts"], sort_keys=True))


if __name__ == "__main__":
    main()
