#!/usr/bin/env python3
"""決定性檢查：同一組條件重複執行，逐位元組比對每一個 token 的 logits。

為什麼需要這個工具
------------------
`sdq/experiments.py` 會用「同一個提示詞跑幾次，第一個生成 token 是否相同」當
正確性指紋。但在 S7/S8 那兩輪，這個指紋偶爾會跳出「⚠️輸出不一致」，
看起來像是 prefill 的累加順序或執行緒交錯造成的非確定性。

實際上那個警告是 **harness 自己的解析問題**：`re.search(r"assistant:\\s*(.*)")`
的 `\\s*` 會跨過換行，assistant 第一行是空的時候就抓到下一行的內容，
而下一行在哪裡換行取決於 stdout 有沒有剛好切成兩個 chunk。

這個工具直接比對 logits（`--dump-logits --logits-out`），把「是不是真的非確定」
用數據回答，而不是靠字串。量測結果（2026-10-04，本機，預設條件）：

    4 次相同執行 → 48 個 token × 248320 個 vocab 的 logits **全部逐位元組相同**

也就是說分頁路徑在同一台機器、同一組條件下是**決定性**的；
設定 `--pager-off` 之外的任何變數才會讓輸出改變。

用法
----
    python tools/determinism_check.py --model <gguf> --bin <epq-chat> -n 4
    python tools/determinism_check.py --model <gguf> --bin <epq-chat> \\
        --runs 3 --env SDQ_FORCE_UNIFORM_SLOT=1     # 比較不同設定是否也一致
"""

from __future__ import annotations

import argparse
import os
import struct
import subprocess
import sys
import tempfile
import time
from pathlib import Path

import numpy as np

REPO = Path(__file__).resolve().parent.parent
DEFAULT_PROMPT = "請用三句話說明什麼是 SSD、RAM、CPU 之間的資料分層。"


def load_logits(path: str) -> np.ndarray:
    """讀 epq-chat --logits-out 寫的檔案：(n_token, vocab) 的 float32。"""
    b = Path(path).read_bytes()
    n = struct.unpack_from("<i", b, 0)[0]
    a = np.frombuffer(b, dtype=np.float32, offset=4)
    k = a.size // n
    return a[: k * n].reshape(k, n)


def run_once(binary: str, model: str, prompt: str, out: str, n_predict: int,
             env_extra: dict, threads: int, ctx: int) -> float:
    env = dict(os.environ)
    env.update(env_extra)
    cmd = [binary, "-m", model, "-p", prompt, "-n", str(n_predict), "-c", str(ctx),
           "-t", str(threads), "--temp", "0", "--no-chat", "--no-interactive",
           "--dump-logits", "--logits-out", out, "--ram-budget-mb", "8192"]
    t0 = time.time()
    p = subprocess.run(cmd, capture_output=True, text=True, env=env)
    if p.returncode != 0:
        print(f"[det] 執行失敗 rc={p.returncode}\n{p.stderr[-1500:]}")
        sys.exit(2)
    return time.time() - t0


def main() -> int:
    ap = argparse.ArgumentParser(description="SDQ 決定性檢查（逐位元組比對 logits）")
    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("--runs", type=int, default=4, help="重複執行幾次")
    ap.add_argument("-n", "--n-predict", type=int, default=16,
                    help="生成幾個 token（每個都會存一份 logits）")
    ap.add_argument("--threads", type=int, default=8)
    ap.add_argument("--ctx", type=int, default=4096)
    ap.add_argument("--prompt", default=DEFAULT_PROMPT)
    ap.add_argument("--env", action="append", default=[],
                    help="額外的環境變數（可重複），用來比較不同設定")
    args = ap.parse_args()

    env_extra = dict(kv.split("=", 1) for kv in args.env)
    print(f"[det] 條件：runs={args.runs} n_predict={args.n_predict} "
          f"threads={args.threads} ctx={args.ctx} env={env_extra or '{}'}")

    with tempfile.TemporaryDirectory() as td:
        mats = []
        for r in range(args.runs):
            out = str(Path(td) / f"logits_{r}.bin")
            dt = run_once(args.bin, args.model, args.prompt, out, args.n_predict,
                          env_extra, args.threads, args.ctx)
            m = load_logits(out)
            mats.append(m)
            print(f"  run {r}: {dt:5.1f}s  logits shape={m.shape}")

        ref = mats[0]
        all_same = True
        for r, m in enumerate(mats[1:], start=1):
            if m.shape != ref.shape:
                print(f"  ❌ run{r} 的 shape 與 run0 不同：{m.shape} vs {ref.shape}")
                all_same = False
                continue
            diff_rows = np.where(np.abs(ref - m).max(axis=1) > 0)[0]
            if diff_rows.size == 0:
                print(f"  ✅ run{r} 與 run0 **逐位元組完全相同**"
                      f"（{ref.shape[0]} token × {ref.shape[1]} vocab）")
            else:
                all_same = False
                i = int(diff_rows[0])
                d = float(np.abs(ref[i] - m[i]).max())
                print(f"  ❌ run{r} 從第 {i} 個 token 開始有差異（max abs={d:.3e}，"
                      f"argmax run0={ref[i].argmax()} run{r}={m[i].argmax()}）")

    print()
    print("[det] 結論：", "決定性 ✅（同條件重複執行逐位元組一致）" if all_same
          else "非決定性 ❌（同條件重複執行會有差異）")
    return 0 if all_same else 1


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