#!/usr/bin/env python3
"""Assemble and validate the Qwen3-235B Frontier best-effort profiles."""

from __future__ import annotations

import argparse
import hashlib
import json
from datetime import datetime, timezone
from pathlib import Path

import pandas as pd


MODEL = "Qwen3-235B-A22B-FP8"
HARDWARE = "h20"
TOKENS = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
QUANT_SIGNATURE = "method=fp8|act=dynamic|serialized=True|block=128x128"
MEASUREMENT_TYPE = "CUDA_EVENT"
PROFILING_PRECISION = "BF16"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--linear-tp4", type=Path, required=True)
    parser.add_argument("--linear-tp8", type=Path, required=True)
    parser.add_argument("--attention-standard", type=Path, required=True)
    parser.add_argument("--attention-mixed", type=Path, required=True)
    parser.add_argument("--moe-tp4-ep1", type=Path, required=True)
    parser.add_argument("--moe-tp1-ep8", type=Path, required=True)
    parser.add_argument("--moe-standalone-tp4-ep1", type=Path, required=True)
    parser.add_argument("--moe-standalone-tp1-ep8", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument(
        "--frontier-commit",
        default="d9cfeb6d8791fbf2f295dd9744c56a666171776e",
    )
    return parser.parse_args()


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


def read_csv(path: Path, label: str) -> pd.DataFrame:
    if not path.is_file():
        raise FileNotFoundError(f"{label}: missing input CSV: {path}")
    data = pd.read_csv(path)
    if data.empty:
        raise ValueError(f"{label}: input CSV is empty: {path}")
    return data


def require_columns(data: pd.DataFrame, columns: list[str], label: str) -> None:
    missing = [column for column in columns if column not in data.columns]
    if missing:
        raise ValueError(f"{label}: missing columns: {missing}")


def require_exact_values(
    data: pd.DataFrame, column: str, expected: set[object], label: str
) -> None:
    actual = set(data[column].dropna().unique().tolist())
    if actual != expected:
        raise ValueError(
            f"{label}: {column} values mismatch: expected={sorted(expected)!r}, "
            f"actual={sorted(actual)!r}"
        )


def validate_common(data: pd.DataFrame, label: str) -> None:
    columns = ["quant_signature", "profiling_precision", "measurement_type"]
    require_columns(data, columns, label)
    require_exact_values(data, "quant_signature", {QUANT_SIGNATURE}, label)
    require_exact_values(data, "profiling_precision", {PROFILING_PRECISION}, label)
    require_exact_values(data, "measurement_type", {MEASUREMENT_TYPE}, label)


def validate_time_columns(
    data: pd.DataFrame, columns: list[str], label: str
) -> None:
    require_columns(data, columns, label)
    if data[columns].isna().any(axis=None):
        counts = data[columns].isna().sum()
        raise ValueError(f"{label}: NaN timing values: {counts[counts > 0].to_dict()}")
    if (data[columns] < 0).any(axis=None):
        raise ValueError(f"{label}: negative timing value")


def validate_token_grid(
    data: pd.DataFrame,
    group_columns: list[str],
    expected_rows_per_point: int,
    label: str,
) -> None:
    for group, rows in data.groupby(group_columns, dropna=False):
        counts = rows["num_tokens"].value_counts().to_dict()
        expected = {token: expected_rows_per_point for token in TOKENS}
        if counts != expected:
            raise ValueError(
                f"{label}: token grid mismatch for group={group}: "
                f"expected={expected}, actual={counts}"
            )


def assemble_linear(tp4: pd.DataFrame, tp8: pd.DataFrame) -> pd.DataFrame:
    label = "linear"
    for source_label, data, expected_tp in (
        ("linear-tp4", tp4, {1, 4}),
        ("linear-tp8", tp8, {1, 8}),
    ):
        validate_common(data, source_label)
        require_columns(data, ["num_tensor_parallel_workers", "num_tokens"], source_label)
        require_exact_values(data, "num_tensor_parallel_workers", expected_tp, source_label)
        validate_token_grid(data, ["num_tensor_parallel_workers"], 1, source_label)

    # The TP=1 rows are replicated operators. Keep the rows from the TP4 run and
    # add only the TP=8 sharded operators from the second run.
    combined = pd.concat(
        [tp4, tp8[tp8["num_tensor_parallel_workers"] == 8]], ignore_index=True
    )
    require_exact_values(combined, "num_tensor_parallel_workers", {1, 4, 8}, label)
    validate_token_grid(combined, ["num_tensor_parallel_workers"], 1, label)

    replicated = combined[combined["num_tensor_parallel_workers"] == 1]
    sharded = combined[combined["num_tensor_parallel_workers"] > 1]
    validate_time_columns(
        replicated,
        [
            "time_stats.emb.median",
            "time_stats.input_layernorm.median",
            "time_stats.post_attention_layernorm.median",
        ],
        "linear replicated operators",
    )
    validate_time_columns(
        sharded,
        [
            "time_stats.attn_pre_proj.median",
            "time_stats.attn_rope.median",
            "time_stats.attn_post_proj.median",
        ],
        "linear sharded attention operators",
    )
    return combined.sort_values(
        ["num_tensor_parallel_workers", "num_tokens"], kind="stable"
    ).reset_index(drop=True)


def assemble_attention(standard: pd.DataFrame, mixed: pd.DataFrame) -> pd.DataFrame:
    timing_columns = [
        "time_stats.attn_input_reshape.median",
        "time_stats.attn_kv_cache_save.median",
        "time_stats.attn_prefill.median",
        "time_stats.attn_decode.median",
        "time_stats.attn_output_reshape.median",
    ]
    for label, data, expected_mixed, expected_rows in (
        ("attention-standard", standard, {False}, 390),
        ("attention-mixed", mixed, {True}, 336),
    ):
        validate_common(data, label)
        require_columns(
            data,
            ["num_tensor_parallel_workers", "is_prefill", "is_mixed_batch"],
            label,
        )
        require_exact_values(data, "num_tensor_parallel_workers", {4, 8}, label)
        require_exact_values(data, "is_prefill", {True}, label)
        require_exact_values(data, "is_mixed_batch", expected_mixed, label)
        validate_time_columns(data, timing_columns, label)
        if len(data) != expected_rows:
            raise ValueError(
                f"{label}: row count mismatch: expected={expected_rows}, actual={len(data)}"
            )

    combined = pd.concat([standard, mixed], ignore_index=True)
    if len(combined) != 726:
        raise ValueError(f"attention: expected 726 rows, got {len(combined)}")
    sort_columns = [
        "num_tensor_parallel_workers",
        "is_mixed_batch",
        "total_tokens",
        "total_prefill_tokens",
        "kv_cache_size",
        "batch_size",
    ]
    return combined.sort_values(sort_columns, kind="stable").reset_index(drop=True)


def assemble_moe(
    tp4_ep1: pd.DataFrame,
    tp1_ep8: pd.DataFrame,
    standalone_tp4_ep1: pd.DataFrame,
    standalone_tp1_ep8: pd.DataFrame,
) -> pd.DataFrame:
    timing_columns = [
        "time_stats.moe_gating_linear.median",
        "time_stats.moe_gating_routing_topk.median",
        "time_stats.moe_shuffling.median",
        "time_stats.moe_grouped_gemm.median",
    ]
    cases = (
        ("moe-tp4-ep1", tp4_ep1, 4, 1, 128, "prefill_hot"),
        ("moe-tp1-ep8", tp1_ep8, 1, 8, 16, "prefill_hot"),
        (
            "moe-standalone-tp4-ep1",
            standalone_tp4_ep1,
            4,
            1,
            128,
            "standalone_legacy",
        ),
        (
            "moe-standalone-tp1-ep8",
            standalone_tp1_ep8,
            1,
            8,
            16,
            "standalone_legacy",
        ),
    )
    for (
        label,
        data,
        expected_tp,
        expected_ep,
        expected_local_experts,
        expected_context,
    ) in cases:
        validate_common(data, label)
        require_columns(
            data,
            [
                "num_tensor_parallel_workers",
                "expert_parallel_size",
                "num_experts_per_device",
                "num_tokens",
                "load_distribution",
                "seed",
                "routing_runtime_path",
                "gating_runtime_context",
            ],
            label,
        )
        require_exact_values(data, "num_tensor_parallel_workers", {expected_tp}, label)
        require_exact_values(data, "expert_parallel_size", {expected_ep}, label)
        require_exact_values(data, "num_experts_per_device", {expected_local_experts}, label)
        require_exact_values(
            data,
            "load_distribution",
            {"uniform", "skewed", "extremely_skewed"},
            label,
        )
        require_exact_values(data, "seed", {0, 1}, label)
        require_exact_values(data, "routing_runtime_path", {"standard_fused_topk"}, label)
        require_exact_values(
            data, "gating_runtime_context", {expected_context}, label
        )
        validate_token_grid(
            data,
            ["num_tensor_parallel_workers", "expert_parallel_size"],
            6,
            label,
        )
        validate_time_columns(
            data,
            timing_columns if expected_context == "prefill_hot" else timing_columns[:2],
            label,
        )
        if len(data) != 90:
            raise ValueError(f"{label}: expected 90 rows, got {len(data)}")

    # Standalone rows are needed only for Frontier's pure-decode gating models.
    # Do not duplicate the shuffling/grouped-GEMM observations merely because
    # the profiler collected them while measuring a second gating context.
    standalone = pd.concat(
        [standalone_tp4_ep1, standalone_tp1_ep8], ignore_index=True
    )
    standalone[timing_columns[2:]] = float("nan")
    combined = pd.concat([tp4_ep1, tp1_ep8, standalone], ignore_index=True)
    return combined.sort_values(
        [
            "num_tensor_parallel_workers",
            "expert_parallel_size",
            "gating_runtime_context",
            "num_tokens",
            "load_distribution",
            "seed",
        ],
        kind="stable",
    ).reset_index(drop=True)


def main() -> None:
    args = parse_args()
    input_paths = {
        "linear_tp4": args.linear_tp4,
        "linear_tp8": args.linear_tp8,
        "attention_standard": args.attention_standard,
        "attention_mixed": args.attention_mixed,
        "moe_tp4_ep1": args.moe_tp4_ep1,
        "moe_tp1_ep8": args.moe_tp1_ep8,
        "moe_standalone_tp4_ep1": args.moe_standalone_tp4_ep1,
        "moe_standalone_tp1_ep8": args.moe_standalone_tp1_ep8,
    }
    inputs = {name: read_csv(path, name) for name, path in input_paths.items()}

    outputs = {
        "linear_op.csv": assemble_linear(inputs["linear_tp4"], inputs["linear_tp8"]),
        "attention.csv": assemble_attention(
            inputs["attention_standard"], inputs["attention_mixed"]
        ),
        "moe.csv": assemble_moe(
            inputs["moe_tp4_ep1"],
            inputs["moe_tp1_ep8"],
            inputs["moe_standalone_tp4_ep1"],
            inputs["moe_standalone_tp1_ep8"],
        ),
    }

    output_dir = args.output_root / "compute" / HARDWARE / MODEL
    output_dir.mkdir(parents=True, exist_ok=True)
    output_paths: dict[str, Path] = {}
    for name, data in outputs.items():
        path = output_dir / name
        data.to_csv(path, index=False)
        output_paths[name] = path

    manifest = {
        "generated_at_utc": datetime.now(timezone.utc).isoformat(),
        "frontier_commit": args.frontier_commit,
        "model": MODEL,
        "hardware": HARDWARE,
        "contract": {
            "quant_signature": QUANT_SIGNATURE,
            "profiling_precision": PROFILING_PRECISION,
            "profiling_precision_semantics": "BF16 kernel output/accumulation; weights and activations are block FP8 W8A8",
            "measurement_type": MEASUREMENT_TYPE,
            "linear_tensor_parallel_sizes": [1, 4, 8],
            "attention_tensor_parallel_sizes": [4, 8],
            "moe_layouts": [
                {"tensor_parallel_size": 4, "expert_parallel_size": 1},
                {"tensor_parallel_size": 1, "expert_parallel_size": 8},
            ],
            "moe_gating_runtime_contexts": [
                "prefill_hot",
                "standalone_legacy",
            ],
        },
        "inputs": {
            name: {
                "path": str(path.resolve()),
                "sha256": sha256(path),
                "rows": len(inputs[name]),
            }
            for name, path in input_paths.items()
        },
        "outputs": {
            name: {
                "path": str(path.resolve()),
                "sha256": sha256(path),
                "rows": len(outputs[name]),
            }
            for name, path in output_paths.items()
        },
    }
    manifest_path = args.output_root / "profile_manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()
