#!/usr/bin/env python3
"""Profile vLLM 0.20's Qwen3-30B unquantized MoE kernel at TP-local shapes."""

from __future__ import annotations

import argparse
import json
import math
import statistics
import subprocess
from pathlib import Path
from typing import Any, Callable

import torch
import vllm


VLLM_VERSION = "0.20.0"
VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1"
HIDDEN_DIM = 2048
INTERMEDIATE_DIM = 768
NUM_EXPERTS = 128
TOP_K = 8


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--vllm-source", type=Path, required=True)
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--tp", type=int, choices=(1, 2, 4), nargs="+", default=[1, 2, 4])
    parser.add_argument("--num-tokens", type=int, nargs="+", default=[8])
    parser.add_argument(
        "--routing-modes",
        choices=("uniform_random_logits", "hotset8"),
        nargs="+",
        default=["uniform_random_logits"],
    )
    parser.add_argument("--warmup-iters", type=int, default=3)
    parser.add_argument("--repeats", type=int, default=5)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--check-reference", action="store_true")
    parser.add_argument(
        "--profile-method",
        choices=("cuda_event", "record_function"),
        default="cuda_event",
    )
    parser.add_argument("--frontier-source", type=Path)
    return parser.parse_args()


def git_head(repo: Path) -> str:
    return subprocess.check_output(
        ["git", "-C", str(repo), "rev-parse", "HEAD"], text=True
    ).strip()


def stats_ms(samples: list[float]) -> dict[str, float]:
    return {
        "min": min(samples),
        "max": max(samples),
        "mean": statistics.fmean(samples),
        "median": statistics.median(samples),
        "std": statistics.pstdev(samples),
    }


def measure_kernel_only_ms(
    fn: Callable[[], torch.Tensor],
    *,
    warmup_iters: int,
    repeats: int,
    trace_root: Path,
    operation_name: str,
    record_function_tracer: type,
) -> tuple[torch.Tensor, dict[str, float]]:
    """Use Frontier's KERNEL_ONLY contract, not a CUDA-event relabel."""
    result = None
    for _ in range(warmup_iters):
        result = fn()
    torch.accelerator.synchronize()

    tracer = record_function_tracer(str(trace_root))
    with tracer:
        for _ in range(repeats):
            with torch.profiler.record_function(f"vidur_{operation_name}"):
                result = fn()
    stats = tracer.get_operation_time_stats()
    if operation_name not in stats:
        raise RuntimeError(f"missing RecordFunctionTracer stats for {operation_name}")
    if result is None:
        raise RuntimeError("kernel-only profiler executed no MoE step")
    return result, {name: float(value) for name, value in stats[operation_name].items()}


def routing_inputs(
    mode: str, num_tokens: int, device: torch.device
) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
    from vllm.model_executor.layers.fused_moe import fused_topk

    if mode == "uniform_random_logits":
        logits = torch.randn(
            (num_tokens, NUM_EXPERTS), device=device, dtype=torch.bfloat16
        )
        hidden_for_topk = torch.empty(
            (num_tokens, HIDDEN_DIM), device=device, dtype=torch.bfloat16
        )
        weights, ids, _ = fused_topk(
            hidden_for_topk,
            logits,
            TOP_K,
            renormalize=True,
        )
    elif mode == "hotset8":
        ids = torch.arange(TOP_K, device=device, dtype=torch.int32).repeat(
            num_tokens, 1
        )
        weights = torch.full(
            (num_tokens, TOP_K),
            1.0 / TOP_K,
            device=device,
            dtype=torch.float32,
        )
    else:
        raise ValueError(mode)

    counts = torch.bincount(ids.flatten().to(torch.int64), minlength=NUM_EXPERTS)
    counts_cpu = counts.cpu().tolist()
    mean_load = num_tokens * TOP_K / NUM_EXPERTS
    variance = sum((count - mean_load) ** 2 for count in counts_cpu) / NUM_EXPERTS
    return weights, ids, {
        "active_experts": sum(count > 0 for count in counts_cpu),
        "min_tokens_per_expert": min(counts_cpu),
        "max_tokens_per_expert": max(counts_cpu),
        "load_cv": math.sqrt(variance) / mean_load if mean_load else 0.0,
        "counts": counts_cpu,
    }


def reference_partial_output(
    hidden: torch.Tensor,
    w13_original: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
) -> torch.Tensor:
    output = torch.zeros_like(hidden)
    for token in range(hidden.shape[0]):
        for route in range(TOP_K):
            expert = int(topk_ids[token, route])
            gate_up = torch.mv(w13_original[expert], hidden[token])
            gate, up = gate_up.chunk(2)
            activated = torch.nn.functional.silu(gate) * up
            expert_output = torch.mv(w2[expert], activated)
            output[token].add_(
                expert_output * topk_weights[token, route].to(expert_output.dtype)
            )
    return output


def main() -> None:
    args = parse_args()
    if vllm.__version__ != VLLM_VERSION:
        raise SystemExit(f"expected vLLM {VLLM_VERSION}, got {vllm.__version__}")
    source_head = git_head(args.vllm_source)
    if source_head != VLLM_COMMIT:
        raise SystemExit(f"expected vLLM source {VLLM_COMMIT}, got {source_head}")
    model_config = json.loads(args.model.joinpath("config.json").read_text())
    expected_model = {
        "hidden_size": HIDDEN_DIM,
        "moe_intermediate_size": INTERMEDIATE_DIM,
        "num_experts": NUM_EXPERTS,
        "num_experts_per_tok": TOP_K,
        "norm_topk_prob": True,
        "torch_dtype": "bfloat16",
    }
    observed_model = {key: model_config.get(key) for key in expected_model}
    if observed_model != expected_model:
        raise SystemExit(
            f"model contract mismatch: expected {expected_model}, got {observed_model}"
        )
    if args.profile_method == "record_function" and args.frontier_source is None:
        raise SystemExit("--frontier-source is required for --profile-method record_function")

    from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
    from vllm.model_executor.layers.fused_moe.activation import MoEActivation
    from vllm.model_executor.layers.fused_moe.config import (
        FUSED_MOE_UNQUANTIZED_CONFIG,
        FusedMoEConfig,
        FusedMoEParallelConfig,
        RoutingMethodType,
    )
    from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
        UnquantizedMoeBackend,
        convert_to_unquantized_kernel_format,
        make_unquantized_moe_kernel,
        select_unquantized_moe_backend,
    )
    from vllm.utils.math_utils import next_power_of_2
    from vllm.v1.worker.workspace import init_workspace_manager

    record_function_tracer = None
    if args.profile_method == "record_function":
        import sys

        sys.path.insert(0, str(args.frontier_source.resolve()))
        from frontier.profiling.utils.record_function_tracer import RecordFunctionTracer

        record_function_tracer = RecordFunctionTracer

    device = torch.device(args.device)
    torch.accelerator.set_device_index(device)
    torch.manual_seed(20260716)
    init_workspace_manager(args.device)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    if args.profile_method == "record_function":
        (args.output.parent / "profiler_traces").mkdir(exist_ok=True)
    max_num_tokens = next_power_of_2(max(args.num_tokens))

    rows: list[dict[str, Any]] = []
    for tp in args.tp:
        parallel = FusedMoEParallelConfig(
            tp_size=tp,
            tp_rank=0,
            pcp_size=1,
            pcp_rank=0,
            dp_size=1,
            dp_rank=0,
            ep_size=1,
            ep_rank=0,
            sp_size=1,
            use_ep=False,
            all2all_backend="allgather_reducescatter",
            enable_eplb=False,
        )
        moe_config = FusedMoEConfig(
            num_experts=NUM_EXPERTS,
            experts_per_token=TOP_K,
            hidden_dim=HIDDEN_DIM,
            intermediate_size_per_partition=INTERMEDIATE_DIM // tp,
            num_local_experts=NUM_EXPERTS,
            num_logical_experts=NUM_EXPERTS,
            activation=MoEActivation.SILU,
            device=device,
            routing_method=RoutingMethodType.Renormalize,
            moe_parallel_config=parallel,
            in_dtype=torch.bfloat16,
            max_num_tokens=max_num_tokens,
        )
        # This process profiles one TP-local weight shard.  Keep the global
        # runtime context single-rank so vLLM does not initialize a collective;
        # the action-conditioned shard size remains explicit in moe_config and
        # the real TP2/TP4 all-reduce is profiled in a separate multi-GPU run.
        vllm_config = VllmConfig(
            parallel_config=ParallelConfig(tensor_parallel_size=1)
        )
        with set_current_vllm_config(vllm_config):
            backend, experts_cls = select_unquantized_moe_backend(moe_config)
            if backend != UnquantizedMoeBackend.FLASHINFER_CUTLASS:
                raise SystemExit(
                    "runtime backend mismatch: expected FlashInfer CUTLASS, "
                    f"got {backend.value} at TP={tp}"
                )
            if experts_cls is None:
                raise SystemExit(f"missing experts class for {backend.value}")

            w13_original = torch.empty(
                (NUM_EXPERTS, 2 * (INTERMEDIATE_DIM // tp), HIDDEN_DIM),
                device=device,
                dtype=torch.bfloat16,
            ).uniform_(-0.01, 0.01)
            w2 = torch.empty(
                (NUM_EXPERTS, HIDDEN_DIM, INTERMEDIATE_DIM // tp),
                device=device,
                dtype=torch.bfloat16,
            ).uniform_(-0.01, 0.01)

            class Layer:
                pass

            layer = Layer()
            layer.moe_config = moe_config
            w13_kernel, w2_kernel = convert_to_unquantized_kernel_format(
                backend,
                layer=layer,
                w13_weight=w13_original,
                w2_weight=w2,
            )
            kernel = make_unquantized_moe_kernel(
                quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,
                moe_config=moe_config,
                backend=backend,
                experts_cls=experts_cls,
            )

            reference_checked = False
            for routing_mode in args.routing_modes:
                for num_tokens in args.num_tokens:
                    hidden = torch.empty(
                        (num_tokens, HIDDEN_DIM),
                        device=device,
                        dtype=torch.bfloat16,
                    ).uniform_(-0.1, 0.1)
                    topk_weights, topk_ids, load = routing_inputs(
                        routing_mode, num_tokens, device
                    )

                    def run_kernel() -> torch.Tensor:
                        return kernel.apply(
                            hidden_states=hidden,
                            w1=w13_kernel,
                            w2=w2_kernel,
                            topk_weights=topk_weights,
                            topk_ids=topk_ids,
                            activation=MoEActivation.SILU,
                            global_num_experts=NUM_EXPERTS,
                            expert_map=None,
                            apply_router_weight_on_input=False,
                        )

                    if args.profile_method == "record_function":
                        output, time_ms = measure_kernel_only_ms(
                            run_kernel,
                            warmup_iters=args.warmup_iters,
                            repeats=args.repeats,
                            trace_root=args.output.parent,
                            operation_name="moe_grouped_gemm",
                            record_function_tracer=record_function_tracer,
                        )
                    else:
                        for _ in range(args.warmup_iters):
                            output = run_kernel()
                        torch.accelerator.synchronize()
                        samples: list[float] = []
                        for _ in range(args.repeats):
                            start = torch.cuda.Event(enable_timing=True)
                            end = torch.cuda.Event(enable_timing=True)
                            start.record()
                            output = run_kernel()
                            end.record()
                            torch.accelerator.synchronize()
                            samples.append(float(start.elapsed_time(end)))
                        time_ms = stats_ms(samples)

                    if output.shape != hidden.shape or not torch.isfinite(output).all():
                        raise SystemExit(
                            f"invalid MoE output TP={tp} M={num_tokens} mode={routing_mode}"
                        )
                    if args.check_reference and not reference_checked:
                        check_tokens = min(2, num_tokens)
                        reference = reference_partial_output(
                            hidden[:check_tokens],
                            w13_original,
                            w2,
                            topk_weights[:check_tokens],
                            topk_ids[:check_tokens],
                        )
                        torch.testing.assert_close(
                            output[:check_tokens], reference, atol=0.03, rtol=0.03
                        )
                        reference_checked = True

                    row = {
                        "tensor_parallel_size": tp,
                        "num_tokens": num_tokens,
                        "routing_mode": routing_mode,
                        "backend": backend.value,
                        "intermediate_size_per_partition": INTERMEDIATE_DIM // tp,
                        "output_is_reduced": kernel.output_is_reduced(),
                        "time_ms": time_ms,
                        "routing_load": load,
                    }
                    rows.append(row)
                    print(
                        json.dumps(
                            {
                                "tp": tp,
                                "num_tokens": num_tokens,
                                "routing_mode": routing_mode,
                                "backend": backend.value,
                                "median_ms": row["time_ms"]["median"],
                            },
                            sort_keys=True,
                        ),
                        flush=True,
                    )

            del kernel, w13_kernel, w2_kernel, w13_original, w2
            torch.accelerator.empty_cache()

    payload = {
        "schema_version": "qwen30_vllm020_moe_raw.v1",
        "environment": {
            "vllm_version": vllm.__version__,
            "vllm_source_commit": source_head,
            "torch_version": torch.__version__,
            "torch_cuda": torch.version.cuda,
            "gpu": torch.cuda.get_device_name(device),
            "model": str(args.model),
            "dtype": "bfloat16",
            "weight_quantization": "none",
            "top_k": TOP_K,
            "norm_topk_prob": True,
            "profile_method": args.profile_method,
        },
        "measurement_scope": (
            "one TP-local weight shard: vLLM modular MoE prepare+FlashInfer "
            "CUTLASS experts+finalize; router linear/top-k and TP all-reduce excluded"
        ),
        "rows": rows,
    }
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


if __name__ == "__main__":
    main()
