#!/usr/bin/env python
# mypy: ignore-errors
"""Sweep a set of DeepSeek V3 MoE benchmark configurations."""

from __future__ import annotations

import argparse
import csv
from pathlib import Path
from types import SimpleNamespace

from scripts.benchmarks.deepseek_v3_moe import (
    DTYPE_MAP,
    benchmark_deepseek_v3,
)

DEFAULT_SWEEP = [
    {
        "batch": 4,
        "seq_len": 1024,
        "hidden_size": 2048,
        "moe_intermediate_size": 4096,
        "n_experts": 64,
        "top_k": 4,
        "groups": 4,
    },
    {
        "batch": 8,
        "seq_len": 2048,
        "hidden_size": 2048,
        "moe_intermediate_size": 4096,
        "n_experts": 64,
        "top_k": 4,
        "groups": 4,
    },
    {
        "batch": 8,
        "seq_len": 2048,
        "hidden_size": 4096,
        "moe_intermediate_size": 8192,
        "n_experts": 128,
        "top_k": 8,
        "groups": 8,
    },
    {
        "batch": 8,
        "seq_len": 2048,
        "hidden_size": 4096,
        "moe_intermediate_size": 8192,
        "n_experts": 256,
        "top_k": 8,
        "groups": 8,
    },
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--dtype",
        choices=DTYPE_MAP.keys(),
        default="bf16",
        help="Computation dtype for all benchmarks",
    )
    parser.add_argument(
        "--device",
        default="auto",
        choices=["auto", "cpu", "cuda"],
        help="Execution device",
    )
    parser.add_argument("--warmup", type=int, default=3, help="Warmup iterations")
    parser.add_argument("--iters", type=int, default=15, help="Benchmark iterations")
    parser.add_argument("--seed", type=int, default=0, help="Random seed")
    parser.add_argument(
        "--group-size",
        type=int,
        default=128,
        help="GROUP_SIZE_M used by the Triton kernel",
    )
    parser.add_argument(
        "--uniform-routing",
        action="store_true",
        help="Force uniform routing for every configuration",
    )
    parser.add_argument(
        "--output",
        type=Path,
        help="Optional CSV file to store results",
    )
    return parser.parse_args()


def make_namespace(base: dict, args: argparse.Namespace) -> SimpleNamespace:
    combined = dict(base)
    combined.update(
        {
            "dtype": args.dtype,
            "device": args.device,
            "warmup": args.warmup,
            "iters": args.iters,
            "seed": args.seed,
            "group_size": args.group_size,
            "uniform_routing": args.uniform_routing,
        }
    )
    return SimpleNamespace(**combined)


def main() -> None:  # pragma: no cover - utility script
    args = parse_args()

    header = (
        "batch",
        "seq_len",
        "hidden_size",
        "moe_intermediate",
        "n_experts",
        "top_k",
        "baseline_ms",
        "patched_ms",
        "speedup",
        "min_tokens",
        "max_tokens",
        "max_diff",
    )
    rows = []

    print(
        f"Running sweep on device={args.device} dtype={args.dtype} uniform_routing={args.uniform_routing}"
    )
    print(
        f"{'batch':>5} {'seq':>5} {'hidden':>7} {'experts':>7} {'topk':>4} {'baseline':>12} {'patched':>12} {'speedup':>8}"
    )

    for cfg in DEFAULT_SWEEP:
        ns = make_namespace(cfg, args)
        result = benchmark_deepseek_v3(ns)
        rows.append(
            (
                cfg["batch"],
                cfg["seq_len"],
                cfg["hidden_size"],
                cfg["moe_intermediate_size"],
                cfg["n_experts"],
                cfg["top_k"],
                result["baseline_ms"],
                result["patched_ms"],
                result["speedup"],
                result["min_tokens"],
                result["max_tokens"],
                result["max_diff"],
            )
        )
        print(
            f"{cfg['batch']:>5} {cfg['seq_len']:>5} {cfg['hidden_size']:>7} {cfg['n_experts']:>7} {cfg['top_k']:>4}"
            f" {result['baseline_ms']:>11.3f} ms {result['patched_ms']:>11.3f} ms {result['speedup']:>7.2f}x"
        )

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        with args.output.open("w", newline="") as fp:
            writer = csv.writer(fp)
            writer.writerow(header)
            writer.writerows(rows)
        print(f"Results written to {args.output}")


if __name__ == "__main__":
    main()