"""
=============================================================================
CS349E · 第 15 讲共享模块：agent_lab.py —— 智能体的工具：反馈的信息量、上下文的账、工具调用期间的 KV
=============================================================================

【这个模块要解决什么问题】
    第 14 讲里，LLM 写内核的能力主要来自一个循环：提出方案 → 调用工具（编译、测试、计时、profiler）→ 读反馈 → 再提出。
    这一讲把这个循环拆开来看，分四部分：

    A. 工具的反馈有多少信息？
       一个"有 bug 的分块矩阵乘法内核"：用 numpy 逐块模拟 GPU 内核的执行，可以注入 7 种真实的 bug
       （K 方向的尾巴没加掩码、循环少算一块、网格向下取整、写越界、少了一次同步、FP16 累加、B 的下标写反）。
       六种"工具"（只说对错的测试、给出错误模式的详细测试、用构造的输入（单位矩阵）测试、跑两次看结果是否一样、
       越界检查 memcheck、竞争检查 racecheck），每种有一个假设的耗时。用信息论算每次工具调用能把"是哪个 bug"的不确定性减少多少比特，
       并比较几种"下一步调用哪个工具"的策略。
    B. 上下文的账：
       真实的编译器错误输出有多长、错误元素的完整清单有多长、摘要有多长；一个多轮会话里上下文怎样增长；
       前缀缓存能省多少 prefill；"只保留最近几次观察"这类截断为什么会让前缀缓存失效。
    C. 工具调用期间，KV cache 怎么办（InferCept 的浪费公式）：
       保留（preserve）、丢掉重算（discard）、换到 CPU（swap）三种做法各浪费多少"显存 × 时间"；
       保留和丢掉的盈亏平衡点；"先保留一段时间、超时再丢"（Continuum 的 TTL）和经典的租雪板问题。
    D. 一个服务模拟器：很多个智能体会话同时跑在一块 H100 上（Llama-3.1-8B），每个会话是"生成 → 等工具 → 再生成"的
       多轮循环；比较几种 KV 处理策略的吞吐和完成时间。

【关于数字】A 的 bug 现象来自真实的 numpy 计算；工具的耗时是假设值（量级参考：nvcc 编译几十秒、compute-sanitizer 让程序
    慢一两个数量级）。C、D 的每步时间用第 5 讲 serving_sim.StepModel 的 roofline 模型（不是实测）；工具耗时的分布是假设。
【运行】python3 agent_lab.py   # 自检（约 1 秒）
【依赖】numpy；../lecture01 的 llm_math、hw_specs；../lecture05 的 serving_sim（StepModel）；../prereq5 的 comm_sim（PCIe 参数）；
    ../prereq4 的 cbuild（只用 find_cc 找 C 编译器）。都是只读导入。
"""

import heapq
import math
import os
import subprocess
import sys
import tempfile
from collections import OrderedDict
from dataclasses import dataclass

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
for sub in ("lecture01", "lecture05", "prereq5", "prereq4"):
    sys.path.insert(0, os.path.join(HERE, "..", sub))

from hw_specs import H100  # noqa: E402
from llm_math import LLAMA31_8B  # noqa: E402
from serving_sim import StepModel  # noqa: E402
from comm_sim import PCIE5  # noqa: E402
from cbuild import find_cc  # noqa: E402

F32 = np.float32

# =============================================================================
# A. 一个有 bug 的分块矩阵乘法内核，和几种"工具"
# =============================================================================
BM = BN = BK = 32          # 块大小（像第 10 讲的共享内存分块）

BUGS = {
    "none": ("没有 bug", "—"),
    "k_tail_unmasked": ("K 方向最后一块没加掩码：K 不是 32 的倍数时读到数组外面", "加载 A、B 时加 k < K 的掩码（越界的读成 0）"),
    "k_skip_last": ("K 循环的边界写错（k0 < K − BK），少算最后一块", "循环条件改成 k0 < K"),
    "grid_floor": ("网格大小向下取整（M // BM），最后不满的块没人算", "网格大小用向上取整 cdiv(M, BM)"),
    "store_overflow": ("写回时没检查 n < N，越界写进下一行的开头", "写回 C 时加 n < N 的掩码"),
    "missing_sync": ("加载下一块之前少了一次 __syncthreads，共享内存有竞争", "在循环里加载之前补上 __syncthreads()"),
    "fp16_accum": ("累加器用了 FP16", "累加器改成 FP32"),
    "b_transposed": ("B 的下标写反了（按 B[n, k] 读）", "按 B[k, n] 读"),
}
BUG_NAMES = tuple(BUGS)

SHAPES = {"nice": (128, 128, 128), "edge": (100, 70, 90), "bigK": (64, 64, 4096)}
SHAPE_NOTES = {"nice": "128³（都是块的整数倍）", "edge": "100 × 70 × 90（都不是块的整数倍）", "bigK": "64 × 64 × 4096（K 很长）"}


def tiled_gemm(A, B, bug="none", run=0):
    """
    模拟一个 GPU 分块矩阵乘法内核（块 32 × 32 × 32）的结果，可以注入一个 bug。

    【参数】A：M×K、B：K×N 的 float32；bug：BUGS 的键；run：第几次运行（决定"未初始化的显存"和竞争的随机性）
    【返回】M×N 的 float32
    【每种 bug 怎么模拟】
        k_tail_unmasked：K 补齐到 32 的倍数，补出来的部分按"行主序内存"读：A 读到下一行的开头，B 读到数组之后的垃圾
        k_skip_last    ：只累加 k0 < K − 32 的块
        grid_floor     ：只有 M // 32 × N // 32 个块；没人写的输出保留"未初始化"的垃圾值
        store_overflow ：最后一列块写 32 列：多出来的列（值为 0，因为加载有掩码）落进下一行的开头；
                         同一个块行里，下一行的开头早已写好（块按行主序启动），于是被 0 覆盖
        missing_sync   ：至少有两个 K 步时，每个块的每个 K 步以 8% 的概率让某个 warp 的 4 行用到"下一块"的 A
        fp16_accum     ：逐个 k 累加，每次加完都舍入到 FP16
        b_transposed   ：把 B 的内存当作 N×K 读
    【边界】这些是"现象级"的模拟：没有真的执行 GPU 指令，但每种 bug 留下的痕迹（哪些元素错、错多少、是否随运行变化、
        是否越界）和真实内核一致。
    """
    M, K = A.shape
    N = B.shape[1]
    g = np.random.default_rng(1000 + run)
    Kp = -(-K // BK) * BK
    if bug == "k_tail_unmasked" and Kp > K:
        fa = np.concatenate([A.ravel(), g.standard_normal(Kp, dtype=F32)])
        Ap = np.stack([fa[i * K:i * K + Kp] for i in range(M)])
        Bp = np.concatenate([B, g.standard_normal((Kp - K, N), dtype=F32)])
        return Ap @ Bp
    if bug == "k_skip_last":
        C = np.zeros((M, N), F32)
        for k0 in range(0, K - BK, BK):
            C += A[:, k0:k0 + BK] @ B[k0:k0 + BK, :]
        return C
    if bug == "fp16_accum":
        acc = np.zeros((M, N), np.float16)
        for k in range(K):
            acc = (acc.astype(F32) + np.outer(A[:, k], B[k, :])).astype(np.float16)
        return acc.astype(F32)
    if bug == "b_transposed":
        return A @ B.ravel().reshape(N, K).T
    C = A @ B
    if bug == "grid_floor":
        out = g.standard_normal((M, N), dtype=F32) * 30          # "未初始化"的输出显存
        mm, nn = (M // BM) * BM, (N // BN) * BN
        out[:mm, :nn] = C[:mm, :nn]
        return out
    if bug == "store_overflow":
        Np = -(-N // BN) * BN
        if Np == N:
            return C
        flat = C.ravel().copy()
        ov = Np - N
        for ti in range(-(-M // BM)):
            for i in range(ti * BM, min(M, ti * BM + BM) - 1):
                flat[(i + 1) * N:(i + 1) * N + ov] = 0.0
        return flat.reshape(M, N)
    if bug == "missing_sync":
        steps = -(-K // BK)
        if steps < 2:
            return C
        C = C.copy()
        for ti in range(-(-M // BM)):
            for tj in range(-(-N // BN)):
                for s in range(steps - 1):
                    if g.random() < 0.08:
                        r0 = ti * BM + 4 * int(g.integers(BM // 4))
                        rows = slice(r0, min(M, r0 + 4))
                        cols = slice(tj * BN, min(N, tj * BN + BN))
                        k0, k1 = s * BK, min(K, (s + 1) * BK)
                        k2 = min(K, k1 + (k1 - k0))
                        if k2 - k1 == k1 - k0:
                            C[rows, cols] += A[rows, k1:k2] @ B[k0:k1, cols] - A[rows, k0:k1] @ B[k0:k1, cols]
        return C
    return C


def make_problem(shape, seed=0):
    """按 SHAPES[shape] 生成 A、B（标准正态）和 FP64 的参考答案。"""
    M, N, K = SHAPES[shape]
    rng = np.random.default_rng(seed)
    A = rng.standard_normal((M, K), dtype=F32)
    B = rng.standard_normal((K, N), dtype=F32)
    return A, B, A.astype(np.float64) @ B.astype(np.float64)


def _ranges(idx):
    """把排好序的下标写成 "0–3, 8, 10–12" 这样的区间（最多 4 段）。"""
    if len(idx) == 0:
        return "无"
    segs, start, prev = [], idx[0], idx[0]
    for v in idx[1:]:
        if v != prev + 1:
            segs.append((start, prev))
            start = v
        prev = v
    segs.append((start, prev))
    s = ", ".join(f"{a}" if a == b else f"{a}–{b}" for a, b in segs[:4])
    return s + (f" …（共 {len(segs)} 段）" if len(segs) > 4 else "")


def summarize(C, ref):
    """
    "详细测试"的输出：哪些元素错了、错多少、错的模式。

    【返回】(文字摘要, 观察值) —— 观察值是一个离散的元组 (是否失败, 模式, 误差量级)，用于信息量的计算。
    【模式】none（全对）、all（> 90% 错）、tail（错的全在最后一个不满的块行 / 块列里）、
        head_cols（错的全在每行开头的几列）、full_cols（错的是若干整列）、scattered（其他）。
    【量级】small：最大误差 < 参考值标准差的 10%；large：否则。
    """
    err = np.abs(C.astype(np.float64) - ref)
    bad = err > 1e-2 + 1e-2 * np.abs(ref)
    M, N = ref.shape
    if not bad.any():
        return "PASS：全部元素在容差内", (False, "none", "-")
    rows = np.flatnonzero(bad.any(axis=1))
    cols = np.flatnonzero(bad.any(axis=0))
    frac = float(bad.mean())
    mag = "small" if err.max() < 0.1 * ref.std() else "large"
    mm, nn = (M // BM) * BM, (N // BN) * BN
    tail = bad[:mm, :nn].sum() == 0 and (mm < M or nn < N)
    if frac > 0.9:
        pat = "all"
    elif tail:
        pat = "tail"
    elif cols.max() < BN and 0 not in rows:
        pat = "head_cols"
    elif bad[:, cols].all():
        pat = "full_cols"
    else:
        pat = "scattered"
    full_rows = np.flatnonzero(bad.mean(axis=1) > 0.9)
    full_cols = np.flatnonzero(bad.mean(axis=0) > 0.9)
    rest = bad.copy()
    rest[full_rows, :] = False
    rest[:, full_cols] = False
    parts = ["几乎所有元素都错"] if (len(full_rows) == M or len(full_cols) == N) else []
    if parts:
        rest[:] = False
    if len(full_rows) and len(full_rows) < M and not parts:
        parts.append(f"几乎整行都错的行 {_ranges(full_rows)}")
    if len(full_cols) and len(full_cols) < N and parts[:1] != ["几乎所有元素都错"]:
        parts.append(f"几乎整列都错的列 {_ranges(full_cols)}")
    if rest.any():
        parts.append(f"其余 {int(rest.sum())} 个错误元素在第 {_ranges(np.flatnonzero(rest.any(axis=1)))} 行")
    text = (f"FAIL：{bad.sum()} / {bad.size} 个元素超出容差（{frac:.1%}）；" + "；".join(parts) +
            f"；最大误差 {err.max():.3g}（参考值的标准差 {ref.std():.3g}）")
    return text, (True, pat, mag)


# 工具：名字 → (假设的耗时 秒, 说明)
TOOLS = {
    "test": (3.0, "运行测试，只报告 PASS / FAIL"),
    "test_verbose": (3.0, "运行测试，报告错误元素的模式和误差大小"),
    "test_identity": (3.0, "用构造的输入（B = 单位矩阵）运行测试，报告错误模式"),
    "run_twice": (5.0, "同样的输入跑两次，比较结果是否逐位相同"),
    "memcheck": (40.0, "compute-sanitizer --tool memcheck：越界读写"),
    "racecheck": (120.0, "compute-sanitizer --tool racecheck：共享内存竞争"),
}


def observe(tool, shape, bug, run=0):
    """
    在 shape 上对带 bug 的内核调用一次工具。

    【返回】(观察值, 给"LLM"看的文字)。观察值是可哈希的离散值，用于贝叶斯更新。
    【memcheck】越界读：k_tail_unmasked 在 K 不是块的整数倍时读出 B 的末尾；越界写：store_overflow 在 N 不是块的整数倍时，
        最后一行的溢出写出了 C 的末尾（中间各行的溢出落在 C 自己的内存里，memcheck 看不出来）。
    【racecheck】missing_sync 在至少两个 K 步时一定报告共享内存的写后读竞争（不管这次结果对不对）。
    """
    _, N, K = SHAPES[shape]
    if tool == "test_identity":
        A, _, _ = make_problem(shape)
        B = np.eye(K, N, dtype=F32)
        return summarize(tiled_gemm(A, B, bug, run), A.astype(np.float64) @ B)[::-1]
    if tool in ("test", "test_verbose", "run_twice"):
        A, B, ref = make_problem(shape)
        C = tiled_gemm(A, B, bug, run)
        if tool == "run_twice":
            same = np.array_equal(C, tiled_gemm(A, B, bug, run + 7919))
            return ("same" if same else "differ"), ("两次结果逐位相同" if same else "两次结果不同（非确定性）")
        text, obs = summarize(C, ref)
        if tool == "test":
            return obs[0], ("FAIL" if obs[0] else "PASS")
        return obs, text
    if tool == "memcheck":
        errs = []
        if bug == "k_tail_unmasked" and K % BK:
            errs.append("Invalid __global__ read（越界读，在 B 的末尾之后）")
        if bug == "store_overflow" and N % BN:
            errs.append("Invalid __global__ write（越界写，在 C 的末尾之后）")
        return tuple(errs), ("；".join(errs) if errs else "没有发现越界访问")
    if tool == "racecheck":
        race = bug == "missing_sync" and -(-K // BK) >= 2
        return race, ("发现共享内存的写后读竞争（RAW hazard）" if race else "没有发现竞争")
    raise ValueError(tool)


ACTIONS = [(t, s) for t in TOOLS for s in SHAPES]


def likelihood_table(n_runs=12):
    """
    P(观察 | bug, 动作)：每个 (工具, 形状) 对每个 bug 运行 n_runs 次（不同的 run），统计观察值的频率。

    【返回】dict：action → {bug: {obs: 概率}}
    【边界】确定性的 bug 每次观察都一样；missing_sync、未初始化显存相关的 bug 会随 run 变化。
    """
    table = {}
    for a in ACTIONS:
        tool, shape = a
        table[a] = {}
        for b in BUG_NAMES:
            cnt = {}
            runs = n_runs if tool in ("test", "test_verbose", "test_identity", "run_twice") else 1
            for r in range(runs):
                o, _ = observe(tool, shape, b, r)
                cnt[o] = cnt.get(o, 0) + 1
            table[a][b] = {o: c / runs for o, c in cnt.items()}
    return table


def entropy(p):
    """香农熵（比特）：H = −Σ p log2 p。"""
    p = np.asarray([v for v in p if v > 0])
    return float(-(p * np.log2(p)).sum())


def posterior(prior, table, action, obs):
    """贝叶斯更新：P(bug | obs) ∝ P(obs | bug) · P(bug)。观察在所有 bug 下都不可能时返回原先验。"""
    post = {b: prior[b] * table[action][b].get(obs, 0.0) for b in prior}
    z = sum(post.values())
    if z == 0:
        return dict(prior)
    return {b: v / z for b, v in post.items()}


def info_gain(prior, table, action):
    """
    一次动作的期望信息增益（比特）= H(先验) − E_obs[H(后验)]。

    【例子】均匀先验下 8 个 bug 是 3 比特；一个只说 PASS / FAIL 的测试最多带来 1 比特。
    """
    h0 = entropy(prior.values())
    p_obs = {}
    for b, pb in prior.items():
        for o, po in table[action][b].items():
            p_obs[o] = p_obs.get(o, 0.0) + pb * po
    eh = 0.0
    for o, po in p_obs.items():
        eh += po * entropy(posterior(prior, table, action, o).values())
    return h0 - eh


def run_policy(policy, true_bug, table, seed=0, max_steps=8, allowed=None):
    """
    让一个"下一步调用哪个工具"的策略去找 bug。

    【参数】policy："greedy_bits"（期望信息增益最大）、"greedy_rate"（期望信息增益 / 耗时最大）、
        或者一个动作列表（按顺序循环使用）；allowed：可用的动作（默认全部）
    【返回】dict(steps=[(动作, 观察文字)], posterior, seconds, found)：found = 最后后验概率最大的就是真 bug 且概率 ≥ 0.99
    【停止】后验最大概率 ≥ 0.99，或用完 max_steps 次，或再调用任何工具都不会带来信息
    """
    rng = np.random.default_rng(seed)
    prior = {b: 1.0 / len(BUG_NAMES) for b in BUG_NAMES}
    allowed = allowed or ACTIONS
    steps, secs = [], 0.0
    for i in range(max_steps):
        if max(prior.values()) >= 0.99:
            break
        if isinstance(policy, list):
            a = policy[i % len(policy)]
        else:
            gains = {a: info_gain(prior, table, a) for a in allowed}
            if max(gains.values()) < 1e-9:
                break
            key = (lambda a: gains[a]) if policy == "greedy_bits" else (lambda a: gains[a] / TOOLS[a[0]][0])
            a = max(allowed, key=key)
        run = int(rng.integers(1_000_000))
        o, text = observe(a[0], a[1], true_bug, run)
        prior = posterior(prior, table, a, o)
        steps.append((a, text))
        secs += TOOLS[a[0]][0]
    best = max(prior, key=prior.get)
    return {"steps": steps, "posterior": prior, "seconds": secs, "found": best == true_bug and prior[best] >= 0.99}


# =============================================================================
# B. 上下文的账
# =============================================================================
def est_tokens(text):
    """
    粗略估计 token 数：英文和代码大约每 4 个字符一个 token，汉字大约每个字一个 token（常见的经验值；不同分词器会差 ±30%）。

    【Python 小知识】'\u4e00' <= ch <= '\u9fff' 判断一个字符是不是常用汉字（CJK 统一表意文字区）。
    """
    cjk = sum(1 for ch in text if "\u4e00" <= ch <= "\u9fff")
    return max(1, round(cjk + (len(text) - cjk) / 4))


BUGGY_C = r"""
#include <stddef.h>
void gemm_tiled(const float *A, const float *B, float *C, int M, int N, int K) {
    const int BM = 32, BN = 32, BK = 32;
    for (int i0 = 0; i0 < M; i0 += BM)
        for (int j0 = 0; j0 < N; j0 += BN) {
            float acc[BM][BN];
            for (int k0 = 0; k0 < K; k0 += BK)
                for (int i = 0; i < BM; i++)
                    for (int j = 0; j < BN; j++)
                        for (int k = 0; k < BK; k++)
                            acc[i][j] += A[(i0 + i) * K + k0 + k] * B[(k0 + k) * N + j0 + j]
            for (int i = 0; i < BM; i++)
                for (int j = 0; j < BN; j++)
                    C[(i0 + i) * N + j0 + j] = acc[i][j];
        }
    return acc;
}
int unused_helper(float *x, size_t n) { for (int i = 0; i < n; i++) x[i] *= 2.0; }
"""


def compile_log(src=BUGGY_C):
    """
    用系统的 C 编译器（-Wall -Wextra）编译一段有错的代码，返回完整的错误输出（没有编译器时返回 None）。

    【用途】看一个"真实的"编译器输出有多长——它会原样进入智能体的上下文，除非工具接口先把它压缩。
    """
    cc = find_cc()
    if cc is None:
        return None
    with tempfile.TemporaryDirectory(prefix="cs349e_l15_") as d:
        path = os.path.join(d, "kernel.c")
        with open(path, "w") as f:
            f.write(src)
        r = subprocess.run([cc, "-Wall", "-Wextra", "-c", path, "-o", os.path.join(d, "kernel.o")],
                           capture_output=True, text=True, timeout=60)
    return r.stderr.replace(d, "<tmp>")


def summarize_compile_log(log):
    """接口层的压缩：只保留每个 error 的那一行（去掉路径、源码回显和 warning），最多 3 条。"""
    errs = [ln.split("kernel.c:", 1)[-1] for ln in log.splitlines() if " error: " in ln]
    more = f"（还有 {len(errs) - 3} 个错误）" if len(errs) > 3 else ""
    return f"{len(errs)} errors: " + " | ".join(errs[:3]) + more


def mismatch_dump(C, ref, limit=None):
    """未经压缩的测试输出：每个错误元素一行 "C[i,j] = got (expected exp)"。"""
    err = np.abs(C.astype(np.float64) - ref)
    bad = np.argwhere(err > 1e-2 + 1e-2 * np.abs(ref))
    lines = [f"C[{i},{j}] = {C[i, j]:.6f} (expected {ref[i, j]:.6f})" for i, j in bad[:limit]]
    return "\n".join(lines)


def session_context(turns, prompt, assistant, tool_tokens, history="full", keep_last=5, placeholder=12):
    """
    一个多轮会话的上下文账。

    【参数】turns：轮数；prompt：系统提示 + 任务的 token 数；assistant：每轮模型输出（思考 + 动作）的 token 数；
        tool_tokens：每轮工具输出的 token 数（数或列表）；history："full"（全部保留）、"last_k"（只保留最近 keep_last 次
        工具输出，更早的换成一个 placeholder 个 token 的占位符，SWE-agent 的做法）
    【返回】dict：context（每轮调用 LLM 时的上下文长度列表）、prefill_nocache（不用前缀缓存时每轮要 prefill 的总 token 数）、
        prefill_cache（用前缀缓存时每轮要 prefill 的 token 数：从"第一个和上一轮不同的位置"开始）
    【为什么截断会让缓存失效】last_k 时，每一轮都会把"刚刚变旧"的那次工具输出换成占位符——它前面的内容不变，
        它和它后面的内容都变了，前缀缓存只能用到它之前。
    """
    tool = tool_tokens if isinstance(tool_tokens, list) else [tool_tokens] * turns
    prev = None
    ctx, pf_no, pf_cache = [], [], []
    for t in range(turns):
        # 第 t 轮调用 LLM 时，上下文 = 提示 + 之前 t 轮的（模型输出 + 工具输出）
        segs = [("p", prompt)]
        for j in range(t):
            segs.append(("a", assistant))
            old = history == "last_k" and j < t - keep_last
            segs.append(("o" if not old else "x", placeholder if old else tool[j]))
        total = sum(n for _, n in segs)
        if prev is None:
            common = 0
        else:
            common = 0
            for (k1, n1), (k2, n2) in zip(prev, segs):
                if k1 != k2 or n1 != n2:
                    break
                common += n1
        ctx.append(total)
        pf_no.append(total)
        pf_cache.append(total - common)
        prev = segs
    return {"context": ctx, "prefill_nocache": pf_no, "prefill_cache": pf_cache}


def prefill_gpu_seconds(new_tokens, ctx_before, model=LLAMA31_8B, chip=H100, mfu=0.5):
    """
    在已有 ctx_before 个 token 的上下文后面 prefill new_tokens 个 token 要多少 GPU 秒（计算受限的估计）。

    【公式】FLOPs = new × 每 token 的线性层运算 + 注意力：第 i 个新 token 看 ctx_before + i 个位置。
    """
    lin = new_tokens * model.linear_flops_per_token()
    att = model.attn_flops_per_token(1) * (new_tokens * ctx_before + new_tokens * (new_tokens + 1) / 2)
    return (lin + att) / (chip.peak_flops("bf16") * mfu)


# =============================================================================
# C. 工具调用期间的 KV：InferCept 的浪费公式
# =============================================================================
STEP = StepModel(LLAMA31_8B, H100)
KV_BYTES = LLAMA31_8B.kv_bytes_per_token(2.0)          # 每个 token 的 KV 字节（BF16）：128 KiB
PCIE_BW = PCIE5.bw * 0.8                               # 有效的 GPU ↔ CPU 带宽（PCIe 5.0 x16 的 80%，假设）


def t_fwd(n):
    """InferCept 的 T_fwd：一次前向处理 n 个 token（单个序列的 prefill / 重算）要多久（第 5 讲的 roofline 模型）。"""
    return STEP.prefill_time([int(n)]) if n > 0 else 0.0


def waste_preserve(C, t_int):
    """InferCept 式 (2)：保留——工具调用的 t_int 秒里，C 个 token 的 KV 白占显存。单位：GB·秒。"""
    return t_int * C * KV_BYTES / 1e9


def waste_discard(C, C_other):
    """InferCept 式 (1)：丢掉重算——重算期间这 C 个 token 的显存 + 其他请求（C_other 个 token）被拖长的那段时间。"""
    return t_fwd(C) * (C + C_other) * KV_BYTES / 1e9


def waste_chunked(C, C_other, n_chunks):
    """InferCept 式 (4)：分块重算——自身那一项减半；其他请求那一项 = 每块多出的迭代时间 × 块数。"""
    return (t_fwd(C) * C / 2 + n_chunks * t_fwd(C / n_chunks) * C_other) * KV_BYTES / 1e9


def waste_swap(C, C_batch, bw=PCIE_BW):
    """InferCept 式 (3)：同步换出 + 换入，期间整批（C_batch 个 token）的显存都在等：2 × T_swap(C) × C_batch。"""
    t_swap = C * KV_BYTES / bw
    return 2 * t_swap * C_batch * KV_BYTES / 1e9


def break_even(C, C_other):
    """保留和丢掉重算的盈亏平衡时间 T* = T_fwd(C)·(C + C_other)/C：工具调用短于 T* 时保留更省。"""
    return t_fwd(C) * (C + C_other) / C


# 一个"写内核的智能体"每轮调用的工具：(名字, 概率, 平均耗时 秒, 对数正态的 σ)。耗时是假设值。
TOOL_MIX = [
    ("读文件 / 搜代码", 0.40, 0.3, 0.5),
    ("编译 + 测试", 0.30, 25.0, 0.3),
    ("只跑测试", 0.15, 4.0, 0.4),
    ("profiler", 0.15, 15.0, 0.4),
]


def sample_tool(rng, mix=None):
    """按 mix（默认 TOOL_MIX）抽一次工具调用：返回 (下标, 耗时)。对数正态的均值等于表里的平均耗时。"""
    mix = mix or TOOL_MIX
    p = np.array([m[1] for m in mix])
    k = int(rng.choice(len(mix), p=p / p.sum()))
    _, _, mean, sig = mix[k]
    return k, float(rng.lognormal(math.log(mean) - sig ** 2 / 2, sig))


# =============================================================================
# D. 很多个智能体会话同时跑：KV 处理策略的服务模拟
# =============================================================================
@dataclass
class AgentWorkload:
    """
    一个智能体会话：系统提示 + 任务 prompt 个 token，然后 turns 轮"生成 decode 个 token → 调用工具 → 工具输出 tool_out 个 token"。

    【字段】agents：同时在跑的会话数（闭环：一个会话结束，马上开始下一个）
    """
    agents: int = 64
    prompt: int = 3000
    turns: int = 12
    decode: int = 150
    tool_out: int = 600


@dataclass
class Session:
    sid: int
    start: float
    ctx: int = 0              # 逻辑上的上下文长度（token 数）
    resident: int = 0         # 在显存里的 KV token 数（钉住的，不能被换出）
    reserved: int = 0         # 为这一轮预留的显存（prefill + decode）
    turn: int = 0
    need: int = 0             # 还要 prefill 的 token 数（包括重算的部分）
    recomp_left: int = 0      # need 里有多少是"重算已有的上下文"（不增加 ctx）
    paused_at: float = 0.0    # 这次工具调用开始的时刻
    decoded: int = 0
    state: str = "queued"     # queued / swapin / prefill / decode / tool / done
    where: str = "none"       # KV 在哪：gpu / cache / cpu / none
    tool_kind: int = -1
    ready_at: float = 0.0     # 这一轮"轮到 LLM"的时刻（统计 LLM 侧的等待）
    recompute: int = 0


POLICIES = ("discard", "prefix_cache", "preserve", "ttl", "ttl_queue", "swap_sync", "swap_async")


class AgentServer:
    """
    一块 GPU 上的推理引擎 + 一群智能体会话，按引擎的迭代推进时间。

    【引擎】每次迭代：所有正在 decode 的会话各生成 1 个 token，再按先来先服务塞进 prefill 的块，总共不超过 token_budget
        （分块 prefill，第 5 讲）；迭代时间 = StepModel.mixed_time(...)。显存按 token 记账，容量 = StepModel.kv_capacity_tokens()。
    【工具调用时（拦截点）怎样处理这个会话的 KV】
        discard     ：立刻释放；工具回来后整个上下文重算（InferCept 说的今天的引擎的默认做法）
        prefix_cache：释放，但留在一个可被挤掉的缓存里（LRU）；回来时如果还在，就不用重算（vLLM 的自动前缀缓存）
        preserve    ：一直钉在显存里
        swap_sync   ：同步换到 CPU：换出和换入的传输时间加到引擎的迭代时间上（所有请求一起等）
        swap_async  ：异步换到 CPU：PCIe 链路是一个先来先服务的队列，传输和计算重叠
        ttl         ：先钉住 TTL 秒（按 InferCept 的浪费公式算的盈亏平衡点 T*），超时还没回来就转成可被挤掉的缓存
        ttl_queue   ：同上，但拦截时如果已经有会话在排队等显存，就不钉住（TTL = 0）——把"排队的代价"也算进来（Continuum 的思路）
    【统计】完成的会话数、会话的完成时间、每轮 LLM 侧的等待、重算的 token 数、被暂停的会话占的显存。
    【参数】tool_mix：工具耗时的分布（默认 TOOL_MIX）；pcie_bw：GPU ↔ CPU 的有效带宽；horizon / warmup：模拟多久、前多少秒不计入统计
    【简化】每个会话一轮里需要的显存（prefill + decode）在开始 prefill 时一次预留；显存不够时，先挤掉可被挤掉的缓存，
        还不够就等（不抢占正在运行的会话）。不模拟前缀在不同会话之间的共享。
    """

    def __init__(self, policy, wl=AgentWorkload(), step=STEP, token_budget=2048, seed=0, horizon=1200.0,
                 warmup=200.0, pcie_bw=PCIE_BW, tool_mix=None):
        assert policy in POLICIES, policy
        self.policy, self.wl, self.step, self.budget = policy, wl, step, token_budget
        self.tool_mix = tool_mix or TOOL_MIX
        self.rng = np.random.default_rng(seed)
        self.horizon, self.warmup, self.bw = horizon, warmup, pcie_bw
        self.cap = step.kv_capacity_tokens()
        self.pinned = 0                       # 钉住的 KV + 预留
        self.cache = OrderedDict()            # sid → token 数（可被挤掉，LRU 顺序）
        self.cache_tokens = 0
        self.t = 0.0
        self.events = []                      # (时间, 序号, 种类, sid)
        self._seq = 0
        self.sessions = {}
        self.queue = []                       # 等 prefill 的会话（先来先服务）
        self.ready = []                       # 换入完成、可以开始 prefill 的会话
        self.link_free = 0.0                  # 异步换入换出：PCIe 链路什么时候空出来
        self.sync_stall = 0.0                 # 同步换入换出：下一次迭代要额外等多久
        self.next_sid = 0
        self.done = []
        self.turn_waits = []
        self.prefill_tokens = 0
        self.recompute_tokens = 0
        self.paused_mem_integral = 0.0        # ∫ 被暂停会话占的显存 dt
        self.paused_pinned = 0
        self.forced_evictions = 0
        self.iters = 0
        for _ in range(wl.agents):
            self._new_session()

    # ---- 小工具 ----
    def _event(self, t, kind, sid):
        heapq.heappush(self.events, (t, self._seq, kind, sid))
        self._seq += 1

    def _free(self):
        return self.cap - self.pinned - self.cache_tokens

    def _ensure(self, n):
        """确保有 n 个 token 的空闲显存：不够就按 LRU 挤掉缓存。返回是否成功。"""
        while self._free() < n and self.cache:
            sid, tok = self.cache.popitem(last=False)
            self.cache_tokens -= tok
            s = self.sessions[sid]
            if s.where == "cache":
                s.where = "none"
                s.resident = 0
        return self._free() >= n

    def _new_session(self):
        s = Session(self.next_sid, self.t, need=self.wl.prompt, ready_at=self.t)
        self.sessions[s.sid] = s
        self.next_sid += 1
        self.queue.append(s.sid)

    # ---- 拦截与恢复 ----
    def _intercept(self, s):
        """这一轮生成完，开始调用工具：按策略处理它的 KV。"""
        k, dur = sample_tool(self.rng, self.tool_mix)
        s.tool_kind, s.state, s.paused_at = k, "tool", self.t
        self._event(self.t + dur, "tool_done", s.sid)
        tok = s.resident
        if self.policy == "discard":
            self.pinned -= tok
            s.resident, s.where = 0, "none"
        elif self.policy == "prefix_cache":
            self._to_cache(s)
        elif self.policy == "preserve":
            self.paused_pinned += tok
        elif self.policy == "swap_sync":
            self.sync_stall += tok * KV_BYTES / self.bw
            self.pinned -= tok
            s.resident, s.where = 0, "cpu"
        elif self.policy == "swap_async":
            start = max(self.t, self.link_free)
            self.link_free = start + tok * KV_BYTES / self.bw
            self.paused_pinned += tok
            self._event(self.link_free, "swapout_done", s.sid)
        elif self.policy in ("ttl", "ttl_queue"):
            if self.policy == "ttl_queue" and self.queue:
                self._to_cache(s)                       # 已经有人在等显存：不钉住
                return
            self.paused_pinned += tok
            ttl = break_even(max(1, s.ctx), max(0, self.pinned - tok))
            self._event(self.t + ttl, "ttl_expire", s.sid)

    def _to_cache(self, s):
        tok = s.resident
        self.pinned -= tok
        if tok:
            self.cache[s.sid] = tok
            self.cache_tokens += tok
            s.where = "cache"
        else:
            s.where = "none"

    def _resume(self, s):
        """工具调用结束：决定这一轮要 prefill 多少（工具输出 + 可能的重算），放进队列。"""
        if s.where == "gpu" and s.resident:
            if self.policy in ("preserve", "ttl", "ttl_queue", "swap_async"):
                self.paused_pinned -= s.resident
            need = self.wl.tool_out
        elif s.where == "cache":
            tok = self.cache.pop(s.sid)
            self.cache_tokens -= tok
            self.pinned += tok
            need = self.wl.tool_out
        elif s.where == "cpu":
            need = self.wl.tool_out          # KV 在 CPU 上：开始 prefill 之前先换回来（在 _admit 里预留显存、启动传输）
        else:
            need = s.ctx + self.wl.tool_out
            s.recomp_left = s.ctx
            s.recompute += s.ctx
            self.recompute_tokens += s.ctx
            s.resident = 0
        s.need, s.state, s.ready_at = need, "queued", self.t
        self.queue.append(s.sid)

    def _handle(self, kind, sid):
        s = self.sessions[sid]
        if kind == "tool_done":
            self._resume(s)
        elif kind == "swapout_done":
            if s.state == "tool" and s.where == "gpu":
                self.pinned -= s.resident
                self.paused_pinned -= s.resident
                s.resident, s.where = 0, "cpu"
        elif kind == "swapin_done":
            s.state = "prefill"
            self.ready.append(s)
        elif kind == "ttl_expire":
            if s.state == "tool" and s.where == "gpu":
                self.paused_pinned -= s.resident
                self._to_cache(s)

    # ---- 主循环 ----
    def _evict_paused(self, want):
        """
        打破僵局：显存全被"暂停中、钉住"的会话占着、又没有任何会话在运行时，按暂停的先后把它们的 KV 丢掉（回来时重算），
        直到放得下 want 个 token。只在"什么都跑不动"时才做——这就是纯 preserve 在显存不够时的样子。
        """
        paused = sorted((s for s in self.sessions.values() if s.state == "tool" and s.where == "gpu" and s.resident),
                        key=lambda s: s.paused_at)
        for s in paused:
            if self._free() >= want:
                break
            self.pinned -= s.resident
            self.paused_pinned -= s.resident
            s.resident, s.where = 0, "none"
            self.forced_evictions += 1
        return self._free() >= want

    def _admit(self, idle):
        """从队首开始，给等 prefill 的会话预留显存（这一轮的 prefill + decode）。先来先服务，显存不够就停。"""
        admitted = []
        while self.queue:
            s = self.sessions[self.queue[0]]
            extra = s.ctx if s.where == "cpu" else 0         # 要换回来的 KV 也要占显存
            want = s.need + self.wl.decode + extra
            ok = self._ensure(want)
            if not ok and idle and not admitted:
                ok = self._evict_paused(want)
            if not ok:
                break
            self.queue.pop(0)
            self.pinned += want
            s.reserved = s.need + self.wl.decode
            if extra:                                        # 换入：显存已经预留好，开始传输
                s.resident = s.ctx
                s.where = "gpu"
                if self.policy == "swap_sync":
                    self.sync_stall += s.ctx * KV_BYTES / self.bw     # 同步：整个引擎等它传完
                else:
                    start = max(self.t, self.link_free)
                    self.link_free = start + s.ctx * KV_BYTES / self.bw
                    s.state = "swapin"
                    self._event(self.link_free, "swapin_done", s.sid)
                    continue
            s.where, s.state = "gpu", "prefill"
            admitted.append(s)
        return admitted

    def run(self):
        prefilling, decoding = [], []
        while self.t < self.horizon:
            while self.events and self.events[0][0] <= self.t:
                _, _, kind, sid = heapq.heappop(self.events)
                self._handle(kind, sid)
            prefilling += self.ready
            self.ready = []
            prefilling += self._admit(idle=not prefilling and not decoding)
            if not prefilling and not decoding:
                if not self.events:
                    break
                nt = self.events[0][0]
                self.paused_mem_integral += self.paused_pinned * (nt - self.t)
                self.t = nt
                continue
            budget = self.budget - len(decoding)
            chunk, attn, finished = 0, 0.0, []
            for s in prefilling:
                if budget <= 0:
                    break
                c = min(s.need, budget)
                r = min(c, s.recomp_left)                       # 先重算已有的上下文，再处理新 token
                attn += LLAMA31_8B.attn_flops_per_token(1) * c * ((s.ctx - s.recomp_left if r else s.ctx) + c / 2)
                s.recomp_left -= r
                s.need -= c
                s.ctx += c - r
                s.resident += c
                s.reserved -= c
                chunk += c
                budget -= c
                if s.need == 0:
                    finished.append(s)
            ctx_sum = sum(s.ctx for s in decoding)
            dt = self.step.mixed_time(len(decoding), ctx_sum, chunk, attn) + self.sync_stall
            self.sync_stall = 0.0
            self.prefill_tokens += chunk
            self.paused_mem_integral += self.paused_pinned * dt
            self.t += dt
            self.iters += 1
            still = []
            for s in decoding:
                s.decoded += 1
                s.ctx += 1
                s.resident += 1
                s.reserved -= 1
                if s.decoded >= self.wl.decode:
                    self.pinned -= s.reserved
                    s.reserved = 0
                    s.decoded = 0
                    s.turn += 1
                    self.turn_waits.append(self.t - s.ready_at)
                    if s.turn >= self.wl.turns:
                        s.state = "done"
                        self.pinned -= s.resident
                        s.resident = 0
                        self.done.append((s.start, self.t, s.recompute))
                        self._new_session()
                    else:
                        self._intercept(s)
                else:
                    still.append(s)
            for s in finished:
                prefilling.remove(s)
                s.state = "decode"
                still.append(s)
            decoding = still
        return self.metrics()

    def metrics(self):
        win = self.horizon - self.warmup
        fin = [(a, b, r) for a, b, r in self.done if a >= self.warmup]
        jct = [b - a for a, b, _ in fin]
        return {"policy": self.policy, "sessions_per_hour": len(fin) / win * 3600 if win > 0 else 0.0,
                "jct_mean": float(np.mean(jct)) if jct else math.nan,
                "turn_wait_mean": float(np.mean(self.turn_waits)) if self.turn_waits else math.nan,
                "recompute_frac": self.recompute_tokens / max(1, self.prefill_tokens),
                "paused_gb": self.paused_mem_integral / max(self.t, 1e-9) * KV_BYTES / 1e9,
                "cap_gb": self.cap * KV_BYTES / 1e9, "iters": self.iters, "forced_evictions": self.forced_evictions}


# =============================================================================
# 自检
# =============================================================================
def _selftest():
    # A：每种 bug 的特征
    A, B, ref = make_problem("nice")
    assert np.allclose(tiled_gemm(A, B, "none"), ref, atol=1e-3)
    for b in ("k_tail_unmasked", "grid_floor", "store_overflow"):        # 尾巴类 bug 在"整齐"的形状上看不出来
        assert summarize(tiled_gemm(A, B, b), ref)[1][0] is False, b
    A2, B2, ref2 = make_problem("edge")
    assert summarize(tiled_gemm(A2, B2, "grid_floor"), ref2)[1][1] == "tail"
    assert summarize(tiled_gemm(A2, B2, "store_overflow"), ref2)[1][1] == "head_cols"
    assert summarize(tiled_gemm(A2, B2, "b_transposed"), ref2)[1][1] == "all"
    table = likelihood_table(n_runs=6)
    prior = {b: 1 / 8 for b in BUG_NAMES}
    assert abs(entropy(prior.values()) - 3.0) < 1e-12
    assert info_gain(prior, table, ("test", "nice")) <= 1.0 + 1e-9
    assert info_gain(prior, table, ("test_verbose", "edge")) > info_gain(prior, table, ("test", "edge"))
    r = run_policy("greedy_bits", "store_overflow", table, seed=0)
    assert r["found"], r["posterior"]
    # B：上下文账
    sc = session_context(10, 2000, 200, 500)
    assert sc["context"][0] == 2000 and sc["prefill_cache"][1] == 700
    lk = session_context(10, 2000, 200, 500, history="last_k", keep_last=3)
    assert sum(lk["prefill_cache"]) > sum(sc["prefill_cache"])
    # C：浪费公式的单调性与盈亏平衡
    C, Co = 8000, 100000
    T = break_even(C, Co)
    assert abs(waste_preserve(C, T) - waste_discard(C, Co)) < 1e-9
    assert waste_chunked(C, Co, 4) <= waste_discard(C, Co) + 1e-12
    # D：一个短的模拟能跑完，显存不超过容量
    m = AgentServer("preserve", AgentWorkload(agents=8, turns=3), horizon=120, warmup=0).run()
    assert m["sessions_per_hour"] > 0
    print("agent_lab 自检通过。")


if __name__ == "__main__":
    print("== agent_lab 自检 ==")
    _selftest()