#!/usr/bin/env python3
"""One-cell gate for the two Qwen235 vLLM 0.20 MoE runtime backends."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import torch


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--frontier-source", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def routing(tokens: int, experts: int = 128, topk: int = 8):
    ids = torch.arange(tokens * topk, device="cuda", dtype=torch.int64)
    ids = (ids % experts).view(tokens, topk)
    weights = torch.full((tokens, topk), 1.0 / topk, device="cuda")
    return weights, ids


def main() -> None:
    args = parse_args()
    sys.path.insert(0, str(args.frontier_source.resolve()))
    from frontier.profiling.moe.moe_vllm_kernel import profile_fused_moe_kernel

    weights, ids = routing(8)
    cells = []
    cells.append(
        {
            "name": "tp4_ep1_triton",
            "stats": profile_fused_moe_kernel(
                num_tokens=8,
                num_experts=128,
                hidden_dim=4096,
                expert_hidden_dim=1536,
                top_k=8,
                topk_weights=weights,
                topk_ids=ids,
                tensor_parallel_size=4,
                use_fp8=True,
                block_shape=[128, 128],
                warmup_steps=1,
                active_steps=2,
            ),
        }
    )
    expert_map = torch.full((128,), -1, device="cuda", dtype=torch.int32)
    expert_map[:16] = torch.arange(16, device="cuda", dtype=torch.int32)
    cells.append(
        {
            "name": "tp1_ep8_flashinfer_cutlass",
            "stats": profile_fused_moe_kernel(
                num_tokens=8,
                num_experts=16,
                hidden_dim=4096,
                expert_hidden_dim=1536,
                top_k=8,
                topk_weights=weights,
                topk_ids=ids,
                tensor_parallel_size=1,
                use_fp8=True,
                block_shape=[128, 128],
                warmup_steps=1,
                active_steps=2,
                global_num_experts=128,
                expert_map=expert_map,
            ),
        }
    )
    payload = {"schema": "qwen235-v020-frontier-moe-smoke-v1", "cells": cells}
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(json.dumps(payload, sort_keys=True))


if __name__ == "__main__":
    main()