#!/usr/bin/env python3
"""Compare patched Frontier MoE decomposition with vLLM's serving path."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch

from frontier.profiling.moe.moe_vllm_kernel import (
    profile_fused_moe_kernel,
    quantize_weights_to_fp8,
)
from vllm.model_executor.layers.fused_moe.fused_moe import (
    fused_experts,
    get_config_dtype_str,
    moe_align_block_size,
    try_get_optimal_moe_config,
)


def _measure(step, warmup: int, active: int) -> dict[str, float]:
    for _ in range(warmup):
        step()
    torch.cuda.synchronize()

    samples = []
    for _ in range(active):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        step()
        end.record()
        torch.cuda.synchronize()
        samples.append(start.elapsed_time(end))
    values = torch.tensor(samples)
    return {
        "min": float(values.min()),
        "median": float(values.median()),
        "mean": float(values.mean()),
        "max": float(values.max()),
        "std": float(values.std()),
    }


def _routing(num_tokens: int, top_k: int, num_experts: int, seed: int):
    generator = torch.Generator(device="cuda")
    generator.manual_seed(seed)
    topk_ids = torch.randint(
        num_experts,
        (num_tokens, top_k),
        generator=generator,
        device="cuda",
        dtype=torch.int64,
    )
    topk_weights = torch.rand(
        (num_tokens, top_k),
        generator=generator,
        device="cuda",
        dtype=torch.float32,
    )
    topk_weights /= topk_weights.sum(dim=-1, keepdim=True)
    return topk_weights, topk_ids


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--tokens", nargs="+", type=int, default=[16, 256, 1024])
    parser.add_argument("--tp", type=int, default=4)
    parser.add_argument("--ep", type=int, default=1)
    parser.add_argument("--warmup", type=int, default=2)
    parser.add_argument("--active", type=int, default=20)
    parser.add_argument("--seed", type=int, default=20260715)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    hidden_dim = 4096
    expert_hidden_dim = 1536
    global_num_experts = 128
    top_k = 8
    block_shape = [128, 128]
    if global_num_experts % args.ep:
        raise ValueError("EP must divide 128 experts")
    if expert_hidden_dim % args.tp:
        raise ValueError("TP must divide the expert intermediate dimension")

    local_num_experts = global_num_experts // args.ep
    local_intermediate = expert_hidden_dim // args.tp
    device = torch.device("cuda")

    torch.manual_seed(args.seed)
    w1_bf16 = torch.randn(
        local_num_experts,
        2 * local_intermediate,
        hidden_dim,
        dtype=torch.bfloat16,
        device=device,
    )
    w2_bf16 = torch.randn(
        local_num_experts,
        hidden_dim,
        local_intermediate,
        dtype=torch.bfloat16,
        device=device,
    )
    w1, w1_scale = quantize_weights_to_fp8(w1_bf16, block_shape=block_shape)
    w2, w2_scale = quantize_weights_to_fp8(w2_bf16, block_shape=block_shape)
    del w1_bf16, w2_bf16
    torch.cuda.empty_cache()

    rows = []
    for index, num_tokens in enumerate(args.tokens):
        topk_weights, topk_ids = _routing(
            num_tokens,
            top_k,
            global_num_experts,
            args.seed + index,
        )
        hidden_states = torch.randn(
            num_tokens,
            hidden_dim,
            dtype=torch.bfloat16,
            device=device,
        )

        frontier_grouped = profile_fused_moe_kernel(
            num_tokens=num_tokens,
            num_experts=local_num_experts,
            hidden_dim=hidden_dim,
            expert_hidden_dim=expert_hidden_dim,
            top_k=top_k,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            tensor_parallel_size=args.tp,
            dtype=torch.bfloat16,
            warmup_steps=args.warmup,
            active_steps=args.active,
            use_fp8=True,
            per_channel_quant=False,
            block_shape=block_shape,
            global_num_experts=global_num_experts,
        )

        config = try_get_optimal_moe_config(
            w1_shape=w1.shape,
            w2_shape=w2.shape,
            top_k=top_k,
            dtype=get_config_dtype_str(
                torch.bfloat16,
                use_fp8_w8a8=True,
            ),
            M=num_tokens,
            block_shape=block_shape,
        )

        def align_step() -> None:
            moe_align_block_size(
                topk_ids,
                config["BLOCK_SIZE_M"],
                global_num_experts,
            )

        alignment = _measure(align_step, args.warmup, args.active)

        def serving_step() -> None:
            fused_experts(
                hidden_states=hidden_states,
                w1=w1,
                w2=w2,
                topk_weights=topk_weights,
                topk_ids=topk_ids,
                inplace=True,
                use_fp8_w8a8=True,
                per_channel_quant=False,
                global_num_experts=global_num_experts,
                w1_scale=w1_scale,
                w2_scale=w2_scale,
                block_shape=block_shape,
            )

        serving = _measure(serving_step, args.warmup, args.active)
        decomposed_ms = frontier_grouped["median"] + alignment["median"]
        rows.append(
            {
                "num_tokens": num_tokens,
                "block_size_m": config["BLOCK_SIZE_M"],
                "frontier_grouped_ms": frontier_grouped,
                "frontier_alignment_ms": alignment,
                "frontier_decomposed_median_ms": decomposed_ms,
                "vllm_fused_experts_ms": serving,
                "decomposed_over_serving": decomposed_ms / serving["median"],
            }
        )

    result = {
        "contract": "Frontier grouped_gemm + shuffling alignment vs vLLM fused_experts",
        "model_shape": {
            "hidden_dim": hidden_dim,
            "expert_hidden_dim": expert_hidden_dim,
            "global_num_experts": global_num_experts,
            "local_num_experts": local_num_experts,
            "top_k": top_k,
            "tp": args.tp,
            "ep": args.ep,
            "dtype": "block_fp8_w8a8_bf16_output",
            "block_shape": block_shape,
        },
        "warmup": args.warmup,
        "active": args.active,
        "rows": rows,
    }
    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))


if __name__ == "__main__":
    main()