#!/usr/bin/env python3
"""SDQ 最佳化實驗 harness：跑一組設定矩陣，產出可重現、可比較的結果。

設計原則
--------
1. **每次量測前丟掉模型檔的 page cache**（POSIX_FADV_DONTNEED），否則上一輪留在
   核心記憶體的頁會讓 SSD 讀取量虛高、命中率虛高。
2. **每個設定跑 N 次取中位數**：單次結果抖動 ±10% 很常見，中位數才有意義。
3. **同時回報 KPI**：decode tok/s、cache 命中率、SSD 讀取量/token、峰值 RSS 與
   本行程 VmSwap。VmSwap=0 僅表示該行程取樣為 0，不代表主機停用了 swap。
4. **自動檢查正確性指紋**：把每個實驗產生的第一個 token 記下來，
   如果某個最佳化讓第一個 token 變了，代表它改變了模型行為 → 必須標記為失敗。
   （速度不能拿正確性換，這是本專案的硬規則。）
5. 結果寫成 JSON + Markdown 表格，方便直接貼進 docs/。

用法
----
    python -m sdq.experiments --model <gguf> --bin <epq-chat> --suite quick
    python -m sdq.experiments --model <gguf> --bin <epq-chat> --suite full --repeat 3
    python -m sdq.experiments ... --only thread,direct      # 只跑某些實驗
"""

from __future__ import annotations

import argparse
import json
import os
import re
import resource
import statistics
import subprocess
import sys
import time
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent

# 固定的量測條件：改這些會讓不同實驗之間無法比較
FIXED = {
    "ctx": "4096",
    # S5 起引擎的 n_ubatch 預設是 128（512 在 8 GB 預算下會讓長提示詞 bad_alloc，
    # 見 docs/05 §5.1）。這裡必須跟著預設走，否則量到的不是我們要交付的那條路徑。
    "ubatch": "128",
    "temp": "0",
    "n_predict": "48",
    "ram_mb": "8192",
}

PROMPTS = [
    "請用三句話說明什麼是 SSD、RAM、CPU 之間的資料分層。",
    "Write a short paragraph about how mixture-of-experts models route tokens.",
    "請列出五種日常會用到電腦的工作，並各用一句話說明。",
]


# ------------------------------------------------------------------ 實驗定義
# 每個實驗 = 一組環境變數 / CLI 參數的覆寫。None 代表沿用預設。
def _suite_quick() -> list[dict]:
    """先摸底：threads、I/O 模式、預取、arena、ubatch。成本低、風險低。"""
    out = []
    for t in (4, 6, 8, 10, 12, 16):
        out.append({"name": f"threads={t}", "args": {"threads": str(t)}})
    out += [
        {"name": "prefetch=off", "args": {"no_prefetch": True}},
        {"name": "fadvise=off", "env": {"SDQ_FADV_DONTNEED": "0"}},
        {"name": "arena=2048MB", "args": {"arena_mb": "2048"}},
        {"name": "arena=3072MB", "args": {"arena_mb": "3072"}},
        {"name": "arena=4608MB", "args": {"arena_mb": "4608"}},
        {"name": "ubatch=128", "args": {"ubatch": "128"}},
        # OPT-1：Strata 式熱門 expert 靜態釘選（實測對本模型無效，保留作為對照）
        # Needs a profile produced by tools/route_overlap.py; point SDQ_PROFILE_IN
        # at your own file. Kept as a control even though it measured as no-op.
        {"name": "hot-pin=2048", "args": {},
         "env": {"SDQ_PROFILE_IN": str(REPO / "bench" / "expert_profile.txt"),
                 "SDQ_HOT_PIN": "2048"}},
        # 排程模式對照
        {"name": "modeB(partial=0)", "args": {}, "env": {"SDQ_PARTIAL_CAP_MB": "0"}},
    ]
    return out


def _suite_io() -> list[dict]:
    """I/O 排程：一個 expert 的 gate/up/down 三段要不要同時讀。

    這是 docs/05 的主線實驗。單一 pread 在本機約 7 ms（1.9 MB / 275 MB/s），
    序列讀三段就是 ~21 ms 的延遲鏈，而且這條鏈的長度與執行緒數無關。
    """
    out = [{"name": "baseline", "args": {}}]
    for n in (4, 8, 12, 16, 24, 32, 48):
        out.append({"name": f"io-threads={n}", "args": {},
                    "env": {"SDQ_IO_THREADS": str(n)}})
    out.append({"name": "io-parallel=off", "args": {},
                "env": {"SDQ_IO_PARALLEL": "0"}})
    return out


def _suite_cache() -> list[dict]:
    """快取策略：這是命中率與 I/O 量的主戰場。"""
    return [
        {"name": "baseline", "args": {}},
        {"name": "hot-pin=1024", "args": {}, "env": {"SDQ_HOT_PIN": "1024"}},
        {"name": "hot-pin=3072", "args": {}, "env": {"SDQ_HOT_PIN": "3072"}},
        {"name": "admit=2nd", "args": {}, "env": {"SDQ_ADMIT_MIN_HITS": "2"}},
        {"name": "modeA(cap=256)", "args": {}, "env": {"SDQ_PARTIAL_CAP_MB": "256"}},
    ]


def _suite_slot() -> list[dict]:
    """槽位配置：大小兩種槽位 vs 全部照最大的配。

    背景：這個模型 37 層的 expert 是 1900544 B，層 34/38/39 是 2039808 B（down 用 Q6_K）。
    舊做法所有槽位都照最大的配 → arena 有 6.8% 是浪費掉的，而 arena 在 8 GB 預算裡
    是最貴的資源（實測 +5% 槽位 ≈ +8~9% tok/s）。
    SDQ_FORCE_UNIFORM_SLOT=1 回到舊行為當對照。
    """
    return [
        {"name": "baseline(大小兩種槽位)", "args": {}},
        {"name": "uniform(全部照大槽位)", "args": {}, "env": {"SDQ_FORCE_UNIFORM_SLOT": "1"}},
    ]


def _suite_arena() -> list[dict]:
    """arena 大小掃描：驗證「槽位數 → tok/s」的斜率，並確認 8 GB 預算內的上限。"""
    out = [{"name": "baseline", "args": {}}]
    for mb in (3072, 3584, 4096):
        out.append({"name": f"arena={mb}MB", "args": {"arena_mb": str(mb)}})
    return out


def _suite_pf() -> list[dict]:
    """跨層預取：layer L 開頭就替 layer L+1 排預取。

    背景：end_token() 是在一個 token 的 40 層「全部算完」之後才排預取的，
    那時沒有任何計算可以重疊 → 實測 prefetch 使用率只有 0~8 %，等於整條路徑是關的。
    這一組量的是「把預取搬到每層開頭」有沒有價值。
    """
    return [
        {"name": "baseline(層開頭預取)", "args": {}, "env": {"SDQ_PREFETCH_LAYER_AHEAD": "1"}},
        {"name": "layer-ahead=off", "args": {}, "env": {"SDQ_PREFETCH_LAYER_AHEAD": "0"}},
        {"name": "prefetch=off", "args": {}, "env": {"SDQ_PREFETCH": "0"}},
        {"name": "pf(queue=256,stale=1)", "args": {},
         "env": {"SDQ_PREFETCH_LAYER_AHEAD": "1", "SDQ_PREFETCH_QUEUE": "256",
                 "SDQ_PREFETCH_STALE_TOKENS": "1"}},
    ]


# The former "spec" suite drove speculative decoding through SDQ_NGRAM_DRAFT, an
# env var that no longer exists anywhere in cpp/. Prompt-lookup n-gram drafting was
# measured at zero accepted tokens on this model (see docs/NEGATIVE_RESULTS.md §4)
# and MTP has its own entry point (cpp/sdq_mtp.cpp -> epq-mtp, see docs/MTP.md),
# so the suite was dead code that would silently no-op rather than measure anything.
SUITES = {"quick": _suite_quick, "io": _suite_io, "cache": _suite_cache,
          "slot": _suite_slot, "arena": _suite_arena, "pf": _suite_pf}


# ------------------------------------------------------------------ 執行
def drop_cache(model: str) -> None:
    subprocess.run([sys.executable, str(REPO / "tools" / "drop_model_cache.py"), model],
                   check=False, capture_output=True)


def _limit(ram_mb: str):
    def f():
        n = int(ram_mb) * 1024 * 1024
        resource.setrlimit(resource.RLIMIT_DATA, (n, n))
    return f


def first_generated_token(stdout: str) -> str:
    """取出「第一個生成 token」的**空白正規化**指紋，用來偵測輸出有沒有改變。

    為什麼要把所有空白都拿掉再取前 24 個字元：
    模型的輸出可能以換行開頭，而 stdout 是被分塊寫出的，所以同一次生成在不同
    執行裡「assistant: 後面第一個非空行」的位置不一定相同（例如一次是
    `assistant: <think>`、另一次是 `assistant:` 換行後才出現 `<think>`）。
    如果指紋取「換行前的原始片段」，**完全相同的輸出**會被判成不一致。

    量測（2026-10-04，本機）：同一提示詞連跑 4 次（含 SDQ_FADV_DONTNEED=1
    並丟掉 20.6 GiB page cache），stdout 長度都是 294 bytes、指紋都是 `<think>`；
    `tools/determinism_check.py` 另外證明 48 token × 248320 vocab 的 logits
    **逐位元組一致** → 分頁路徑是決定性的，過去那些 ⚠️ 都是假警報。

    抓不到內容時回空字串，呼叫端不去比對空指紋（寧可少檢查也不要假警報）。
    """
    lines = stdout.splitlines()
    started = False
    chunks: list[str] = []
    for line in lines:
        if not started:
            if line.startswith("assistant:"):
                started = True
                chunks.append(line[len("assistant:"):])
            continue
        if line.startswith("[sdq]"):
            continue
        chunks.append(line)
    text = re.sub(r"\s+", "", "".join(chunks))
    return text[:24]


def run_once(binary: str, model: str, exp: dict, prompt: str, idx: int) -> dict:
    stats_file = Path("/tmp") / f"sdqexp_{os.getpid()}_{idx}.jsonl"
    if stats_file.exists():
        stats_file.unlink()

    cmd = [binary, "-m", model, "-p", prompt, "-n", FIXED["n_predict"],
           "-c", FIXED["ctx"], "-t", str(exp.get("args", {}).get("threads", 8)),
           "-b", str(exp.get("args", {}).get("ubatch", FIXED["ubatch"])), "--temp", FIXED["temp"],
           "--no-chat", "--no-interactive", "--stats-file", str(stats_file),
           "--ram-budget-mb", FIXED["ram_mb"]]
    if exp.get("args", {}).get("arena_mb"):
        cmd += ["--arena-mb", str(exp["args"]["arena_mb"])]
    if exp.get("args", {}).get("no_prefetch"):
        cmd += ["--no-prefetch"]

    env = dict(os.environ)
    env.update(exp.get("env", {}))
    env.setdefault("SDQ_FADV_DONTNEED", "1")

    drop_cache(model)
    t0 = time.time()
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, env=env,
                              preexec_fn=_limit(FIXED["ram_mb"]), timeout=900)
    except subprocess.TimeoutExpired:
        return {"error": "timeout"}
    wall = time.time() - t0
    if proc.returncode != 0:
        return {"error": f"rc={proc.returncode}", "stderr": proc.stderr[-800:]}

    stats = {}
    if stats_file.exists():
        for line in stats_file.read_text().splitlines():
            rec = json.loads(line)
            stats[rec["phase"]] = rec
    stats_file.unlink(missing_ok=True)

    dec = stats.get("decode") or {}
    tok_s = 0.0
    m = re.search(r"decode: (\d+) tok / ([\d.]+)s = ([\d.]+) tok/s", proc.stderr)
    if m:
        tok_s = float(m.group(3))
    peak_rss = 0
    vmswap = 0
    m = re.search(r"峰值 (\d+) MB / (?:swap|本行程 VmSwap) (\d+) MB", proc.stderr)
    if m:
        peak_rss, vmswap = int(m.group(1)), int(m.group(2))

    # 正確性指紋：第一個生成 token（用來偵測「最佳化偷偷改變了模型行為」）
    first_tok = first_generated_token(proc.stdout)

    return {
        "tok_s": tok_s,
        "wall_s": round(wall, 2),
        "hit_rate": dec.get("hit_rate"),
        "ssd_mb_per_tok": (dec.get("ssd_bytes_per_token") or 0) / 2 ** 20,
        "peak_rss_mb": peak_rss,
        "vmswap_mb": vmswap,
        "n_slots": dec.get("n_slots"),
        "evictions": dec.get("evictions"),
        "prefetch_used_rate": dec.get("prefetch_used_rate"),
        "first_tok": first_tok,
        "text": proc.stdout.strip()[-260:],
    }


def run_exp(binary: str, model: str, exp: dict, repeat: int, prompts: list[str]) -> dict:
    runs = []
    for r in range(repeat):
        res = run_once(binary, model, exp, prompts[r % len(prompts)], r)
        if res.get("error"):
            return {"name": exp["name"], "error": res["error"], "stderr": res.get("stderr", "")}
        runs.append(res)
    ok = [r for r in runs if r["tok_s"] > 0]
    if not ok:
        return {"name": exp["name"], "error": "沒有任何一次量到 tok/s"}
    med = lambda k: statistics.median([r[k] for r in ok if r.get(k) is not None]) if any(r.get(k) is not None for r in ok) else None
    # 只比較「同一個提示詞」的重複執行：實驗會輪替提示詞，跨提示詞比字串沒有意義。
    # **空指紋不參與比較**：抓不到第一個 token 時回 ""，那是 stdout 分塊造成的，
    # 不是輸出不一致（見 first_generated_token 的說明）。寧可少檢查也不要假警報。
    by_prompt: dict[str, set] = {}
    for i, r in enumerate(ok):
        if not r["first_tok"]:
            continue
        by_prompt.setdefault(PROMPTS[i % len(PROMPTS)], set()).add(r["first_tok"])
    unstable = [p[:20] for p, v in by_prompt.items() if len(v) > 1]
    # 有沒有真的比到東西（沒比到就不能宣稱「輸出穩定」）
    checked = sum(len(v) for v in by_prompt.values()) - len(by_prompt)
    return {
        "name": exp["name"],
        "args": exp.get("args", {}), "env": exp.get("env", {}),
        "tok_s": round(med("tok_s"), 3),
        "tok_s_all": [r["tok_s"] for r in ok],
        "hit_rate": round(med("hit_rate"), 2) if med("hit_rate") is not None else None,
        "ssd_mb_per_tok": round(med("ssd_mb_per_tok"), 1) if med("ssd_mb_per_tok") is not None else None,
        "peak_rss_mb": med("peak_rss_mb"),
        "vmswap_mb": med("vmswap_mb"),
        "n_slots": med("n_slots"),
        "evictions": med("evictions"),
        "prefetch_used_rate": med("prefetch_used_rate"),
        "stable_output": not unstable,
        "stable_checked": checked,   # 實際有幾組「同提示詞重複」可比
        "unstable_prompts": unstable,
        "first_tok": ok[0]["first_tok"],
        "runs": len(ok),
    }


# ------------------------------------------------------------------ 報告
def to_markdown(results: list[dict]) -> str:
    head = ("| 實驗 | decode tok/s | 命中率 | SSD MB/token | 峰值 RSS | swap | 槽數 | 淘汰 | 預取使用率 | 輸出穩定 |")
    sep = "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | :---: |"
    lines = [head, sep]
    base = next((r for r in results if r["name"] in ("baseline", "threads=8")), None)
    bt = base["tok_s"] if base else None
    for r in results:
        if r.get("error"):
            lines.append(f"| {r['name']} | ❌ {r['error']} | | | | | | | | |")
            continue
        delta = ""
        if bt:
            d = (r["tok_s"] - bt) / bt * 100
            delta = f" ({d:+.1f}%)"
        lines.append(
            f"| {r['name']} | {r['tok_s']}{delta} | {r['hit_rate']}% | {r['ssd_mb_per_tok']} | "
            f"{r['peak_rss_mb']} MB | {r['vmswap_mb']} MB | {r['n_slots']} | {r['evictions']} | "
            f"{r['prefetch_used_rate']}% | "
            f"{'✅' if r['stable_output'] else ('⚠️不一致' if r.get('stable_checked') else '—未比到')} |")
    return "\n".join(lines)


def main() -> int:
    ap = argparse.ArgumentParser(description="SDQ 最佳化實驗")
    ap.add_argument("--model", required=True)
    ap.add_argument("--bin", default=str(REPO.parent / "llama.cpp" / "build" / "bin" / "epq-chat"),
                    help="epq-chat 路徑（預設 <repo>/../llama.cpp/build/bin/epq-chat）")
    ap.add_argument("--suite", default="quick", choices=sorted(SUITES) + ["all"])
    ap.add_argument("--only", default="", help="只跑名稱包含這些關鍵字的實驗（逗號分隔）")
    ap.add_argument("--repeat", type=int, default=2)
    ap.add_argument("--out", default=str(REPO / "bench" / "results"))
    ap.add_argument("--label", default="")
    args = ap.parse_args()

    exps: list[dict] = []
    names = sorted(SUITES) if args.suite == "all" else [args.suite]
    for n in names:
        exps += SUITES[n]()
    # 去掉重複名稱（不同 suite 可能有 baseline）
    seen, uniq = set(), []
    for e in exps:
        if e["name"] in seen:
            continue
        seen.add(e["name"])
        uniq.append(e)
    if args.only:
        keys = [k.strip() for k in args.only.split(",") if k.strip()]
        uniq = [e for e in uniq if any(k in e["name"] for k in keys)]

    print(f"[exp] 模型：{args.model}")
    print(f"[exp] 執行檔：{args.bin}")
    print(f"[exp] 實驗數：{len(uniq)}，每個跑 {args.repeat} 次取中位數")
    print(f"[exp] 固定條件：{FIXED}")
    print()

    results = []
    for e in uniq:
        print(f"[exp] → {e['name']} ...", end=" ", flush=True)
        r = run_exp(args.bin, args.model, e, args.repeat, PROMPTS)
        results.append(r)
        if r.get("error"):
            print(f"❌ {r['error']}")
        else:
            print(f"{r['tok_s']} tok/s  命中率 {r['hit_rate']}%  SSD {r['ssd_mb_per_tok']} MB/tok"
                  + ("" if r["stable_output"] or not r.get("stable_checked")
                     else "  ⚠️ 輸出不一致"))

    Path(args.out).mkdir(parents=True, exist_ok=True)
    label = args.label or args.suite
    jf = Path(args.out) / f"exp_{label}.json"
    jf.write_text(json.dumps({"fixed": FIXED, "repeat": args.repeat,
                              "results": results}, ensure_ascii=False, indent=1), encoding="utf-8")
    md = to_markdown(results)
    mf = Path(args.out) / f"exp_{label}.md"
    mf.write_text(md + "\n", encoding="utf-8")

    print()
    print(md)
    print()
    print(f"[exp] JSON → {jf}")
    print(f"[exp] Markdown → {mf}")
    return 0


if __name__ == "__main__":
    sys.exit(main())