"""Measure what changing the number of active MoE experts actually does.

Run against any MoE GGUF. It loads the model once per expert count, times a
fixed generation, and reports throughput. The point is to replace a guess with
a number before touching a serving configuration.
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

from llama_cpp import Llama

PROMPT = "Write a Python function that reverses a string."


def measure(
    model_path: str,
    expert_count: int | None,
    tokens: int,
    threads: int,
) -> dict:
    """Load the model with an optional expert override and time a generation."""
    key = "expert_used_count"
    overrides = None
    if expert_count is not None:
        overrides = {key: expert_count}

    kwargs = {"kv_overrides": overrides} if overrides else {}
    started = time.monotonic()
    try:
        model = Llama(
            model_path=model_path,
            n_ctx=512,
            n_threads=threads,
            verbose=False,
            **kwargs,
        )
    except Exception as exc:
        return {"experts": expert_count, "error": f"{type(exc).__name__}: {exc}"}
    load_s = time.monotonic() - started

    arch = model.metadata.get("general.architecture", "")
    native = model.metadata.get(f"{arch}.{key}")

    started = time.monotonic()
    out = model.create_chat_completion(
        messages=[{"role": "user", "content": PROMPT}],
        max_tokens=tokens,
        temperature=0.0,
    )
    gen_s = time.monotonic() - started
    produced = out.get("usage", {}).get("completion_tokens", 0)
    text = out["choices"][0]["message"]["content"]
    del model

    return {
        "experts": expert_count,
        "native_experts": native,
        "load_s": round(load_s, 2),
        "gen_s": round(gen_s, 2),
        "tokens": produced,
        "tok_per_s": round(produced / gen_s, 1) if gen_s else None,
        "preview": text[:120],
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("model", help="path to a MoE GGUF file")
    parser.add_argument(
        "--experts",
        default="2,8,15",
        help="comma-separated expert counts to try (omit for native only)",
    )
    parser.add_argument("--tokens", type=int, default=128)
    parser.add_argument("--threads", type=int, default=4)
    parser.add_argument("--json", help="write the results to this JSON file")
    args = parser.parse_args()

    if not Path(args.model).is_file():
        parser.error(f"model not found: {args.model}")

    counts: list[int | None] = [None]
    counts += [int(c) for c in args.experts.split(",") if c.strip()]

    results = []
    for count in counts:
        label = "native" if count is None else f"{count} experts"
        print(f"--- {label}", flush=True)
        row = measure(args.model, count, args.tokens, args.threads)
        results.append(row)
        if "error" in row:
            print(f"    FAILED: {row['error']}", flush=True)
        else:
            print(
                f"    {row['tok_per_s']} tok/s "
                f"(load {row['load_s']}s, gen {row['gen_s']}s)",
                flush=True,
            )

    baseline = next((r for r in results if r.get("experts") is None), None)
    if baseline and baseline.get("tok_per_s"):
        print("\nrelative to native:", flush=True)
        for row in results:
            if row.get("tok_per_s"):
                ratio = row["tok_per_s"] / baseline["tok_per_s"]
                print(
                    f"  {str(row['experts']):>4} experts -> x{ratio:.2f} speed",
                    flush=True,
                )

    if args.json:
        Path(args.json).write_text(
            json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8"
        )
        print(f"\nwritten to {args.json}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())