#!/usr/bin/env python3
"""從 SDQ_ROUTE_LOG 算出「連續 k 個 token 需要幾個不同的 expert」。

這是決定「多 token 批次 / 投機解碼值不值得做」的關鍵數字。
每一層每 token 只會用到 8 個 expert，所以如果把 k 個 token 放在同一層一起算，
需要從 SSD 讀的 expert 數是 U(k) 而不是 8k。U(k)/k 越小，省下的 I/O 越多。

另外會順便算：
  * 完美（Belady）快取命中率 —— 也就是「這個工作流在容量無限時的命中率天花板」
  * 快取容量對命中的影響     —— 用來判斷「把 arena 加大」值不值得
  * 只用過一次的 expert 占比 —— 用來判斷「cache bypass / 准入策略」值不值得

用法
----
SDQ_ROUTE_LOG=/tmp/route.log <epq-chat> -m model.gguf -p "..." -n 48 --no-interactive
    python tools/route_overlap.py /tmp/route.log
"""

from __future__ import annotations

import sys
from collections import defaultdict


def load(path: str) -> dict[int, list[tuple[int, ...]]]:
    """回傳 {layer: [每個 decode step 的 8 個 expert]}（只取 n_tok == 1 的呼叫）。"""
    per_layer: dict[int, list[tuple[int, ...]]] = defaultdict(list)
    with open(path, encoding="utf-8") as f:
        for line in f:
            parts = line.split()
            if len(parts) < 3:
                continue
            il, n_tok = int(parts[0]), int(parts[1])
            if n_tok != 1:
                continue
            per_layer[il].append(tuple(int(x) for x in parts[2].split(":")[1].split(",")))
    return per_layer


def hit_rate_lru(stream, cap: int) -> float:
    cache: dict[int, None] = {}
    hit = tot = 0
    for req in stream:
        for e in req:
            tot += 1
            if e in cache:
                hit += 1
                cache.pop(e)
            elif len(cache) >= cap:
                cache.pop(next(iter(cache)))
            cache[e] = None
    return hit / max(1, tot)


def hit_rate_belady(stream, cap: int) -> float:
    """離線模擬：替換掉「下一次使用最晚」的頁 = 理論上界。"""
    nxt: dict[int, int] = {}
    for i in range(len(stream) - 1, -1, -1):
        for e in stream[i]:
            nxt[e] = i
    cache: set[int] = set()
    hit = tot = 0
    for i, req in enumerate(stream):
        for e in req:
            tot += 1
            if e in cache:
                hit += 1
            else:
                if len(cache) >= cap:
                    cache.discard(max(cache, key=lambda x: nxt.get(x, 10 ** 9)))
                cache.add(e)
    return hit / max(1, tot)


def main() -> int:
    if len(sys.argv) < 2:
        print(__doc__)
        return 2
    per_layer = load(sys.argv[1])
    if not per_layer:
        print("route log 裡沒有 n_tok==1 的紀錄")
        return 1
    n = min(len(v) for v in per_layer.values())
    nl = len(per_layer)
    print(f"每層 decode steps = {n}，層數 = {nl}")
    print()
    print("=== 多 token 批次的 I/O 節省空間 ===")
    print("（分頁一次 = 1.9 MB；k 個 token 分開算要讀 8k 次，合併算只要 U(k) 次）")
    print(f"{'k':>3} {'U(k) 平均':>10} {'U(k)/k':>8} {'省下的讀取':>12}")
    for k in (1, 2, 3, 4, 5, 6, 8):
        tot = cnt = 0
        for v in per_layer.values():
            s = v[:n]
            for i in range(n - k + 1):
                u: set[int] = set()
                for j in range(k):
                    u |= set(s[i + j])
                tot += len(u)
                cnt += 1
        u_avg = tot / cnt
        save = (1 - u_avg / (8.0 * k)) * 100 if k > 1 else 0.0
        print(f"{k:3d} {u_avg:10.2f} {u_avg / k:8.2f} {save:11.1f}%")
    print()
    print("讀法：U(k)/k < 8 代表把 k 個 token 合併到同一層有省到 I/O；")
    print("     合併的 token 越多（k 越大）省得越多，但要能『同時知道』那 k 個 token。")
    print()
    print("=== 快取容量 vs 命中率（每層獨立計算，理想化）===")
    print(f"{'每層槽數':>9} {'LRU':>8} {'Belady(上界)':>13}")
    for cap in (256, 512, 1024, 2048, 4096):
        a = b = 0.0
        for v in per_layer.values():
            s = v[:n]
            a += hit_rate_lru(s, cap)
            b += hit_rate_belady(s, cap)
        print(f"{cap:9d} {a / nl * 100:7.1f}% {b / nl * 100:12.1f}%")
    print()
    print("讀法：LRU == Belady 代表『策略已經最optimal，剩下只有容量問題』；")
    print("     LRU 隨容量上升代表加大 arena 有用。")
    print()
    print("=== 重複使用次數分布（48 個 token 的視窗內）===")
    hist: dict[int, int] = defaultdict(int)
    for v in per_layer.values():
        cnt: dict[int, int] = defaultdict(int)
        for req in v[:n]:
            for e in req:
                cnt[e] += 1
        for e, c in cnt.items():
            hist[c] += c
    total = sum(hist.values())
    once = hist.get(1, 0)
    print(f"只出現 1 次的請求：{once} / {total}（{once / total * 100:.1f}%）")
    print("（占比低 = 快取被單次使用的 expert 污染的機會少，cache-bypass 不值得做）")
    return 0


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