#!/usr/bin/env python3
"""Collect per-layer expert routing histograms from the hybrid GLM-5.2.

Uses vLLM's enable_return_routed_experts to get [seq_len, n_moe_layers,
topk] expert ids per request, over a mixed local calibration corpus
(code + docs + synthetic instructions). Prefill-only (max_tokens=1).

Output: /data/glm52-expert-stats.npz with counts [n_layers, 256].
"""
import glob
import json
import os
import random

import numpy as np

MODEL = "/data/glm52"
OUT = "/data/glm52-expert-stats.npz"
N_EXPERTS = 256
PROMPT_TOKENS = 6000  # approx, chars/4
N_PROMPTS = 192


def build_corpus() -> list[str]:
    random.seed(7)
    texts = []
    # code: vllm sources
    py = sorted(glob.glob("/home/coder/git/glm52/vllm/vllm/**/*.py", recursive=True))
    random.shuffle(py)
    # prose/tech: docs
    md = sorted(glob.glob("/home/coder/git/glm52/vllm/docs/**/*.md", recursive=True))
    random.shuffle(md)

    def chunks(paths, n):
        out, buf = [], ""
        for p in paths:
            try:
                buf += open(p, errors="ignore").read() + "\n\n"
            except OSError:
                continue
            while len(buf) >= PROMPT_TOKENS * 4:
                out.append(buf[: PROMPT_TOKENS * 4])
                buf = buf[PROMPT_TOKENS * 4:]
                if len(out) >= n:
                    return out
        if buf:
            out.append(buf)
        return out[:n]

    texts += chunks(py, N_PROMPTS // 2)          # 50% code
    texts += chunks(md, N_PROMPTS // 4)          # 25% docs
    # 25% synthetic instruction/chat/reasoning
    topics = ["quantum computing", "the French Revolution", "sourdough bread",
              "distributed databases", "protein folding", "jazz harmony",
              "supply chains", "volcanoes", "Rust lifetimes", "photosynthesis",
              "medieval trade routes", "black holes"]
    templates = [
        "Explain {t} to a beginner, covering the key concepts step by step, "
        "common misconceptions, practical examples, and finally an advanced "
        "summary with open research questions. Be very detailed.\n\n",
        "Write a detailed technical design document about building a system "
        "related to {t}: requirements, architecture, tradeoffs, testing plan, "
        "rollout strategy, and failure modes.\n\n",
        "You are a helpful assistant. The user asks a long multi-part "
        "question about {t}. Answer each part with careful reasoning:\n"
        "1) history 2) fundamentals 3) state of the art 4) critiques "
        "5) future directions.\n\n",
    ]
    synth = []
    while len(synth) < N_PROMPTS // 4:
        t = random.choice(topics)
        body = random.choice(templates).format(t=t)
        synth.append((body * 40)[: PROMPT_TOKENS * 4])
    texts += synth
    random.shuffle(texts)
    return texts


def main():
    os.environ.setdefault("VLLM_PP_LAYER_PARTITION", "19,20,21,18")
    stats_dir = "/data/glm52-expert-stats"
    os.environ["VLLM_HYBRID_EXPERT_STATS"] = stats_dir
    from vllm import LLM, SamplingParams

    llm = LLM(
        model=MODEL,
        pipeline_parallel_size=4,
        gpu_memory_utilization=0.51,
        kv_cache_dtype="fp8_ds_mla",
        max_model_len=8192,
        max_num_seqs=2,
        max_num_batched_tokens=4096,
        enforce_eager=True,
    )
    from vllm.inputs import TokensPrompt

    texts = build_corpus()
    tok = llm.get_tokenizer()
    prompts = [
        TokensPrompt(prompt_token_ids=tok.encode(t)[:7900]) for t in texts
    ]
    print(f"{len(prompts)} calibration prompts")
    sp = SamplingParams(max_tokens=1, temperature=0.0)
    outs = llm.generate(prompts, sp)
    total_tokens = sum(len(o.prompt_token_ids) for o in outs)
    del llm

    # AqlmMoEMethod on each PP worker dumped layer_N.npy files
    files = sorted(glob.glob(os.path.join(stats_dir, "layer_*.npy")))
    assert files, f"no stats dumped in {stats_dir}"
    layer_ids = sorted(int(os.path.basename(f)[6:-4]) for f in files)
    counts = np.stack(
        [np.load(os.path.join(stats_dir, f"layer_{li}.npy")) for li in layer_ids]
    )
    np.savez(OUT, counts=counts, layer_ids=np.array(layer_ids),
             total_tokens=total_tokens)
    # quick skew report
    frac = counts / counts.sum(axis=1, keepdims=True).clip(min=1)
    sorted_frac = np.sort(frac, axis=1)[:, ::-1]
    top32 = sorted_frac[:, :32].sum(axis=1)
    top64 = sorted_frac[:, :64].sum(axis=1)
    print(f"tokens: {total_tokens}, layers: {counts.shape[0]}")
    print(f"routing mass in top-32 experts: mean {top32.mean():.3f} "
          f"min {top32.min():.3f} max {top32.max():.3f}")
    print(f"routing mass in top-64 experts: mean {top64.mean():.3f} "
          f"min {top64.min():.3f} max {top64.max():.3f}")
    print("saved", OUT)


if __name__ == "__main__":
    main()