"""
=============================================================================
CS349E · 第 14 讲共享模块：evolve.py —— 迷你 AlphaEvolve：进化一段"选分块配置"的代码
=============================================================================

【这个模块要解决什么问题】
    AlphaEvolve（Novikov et al. 2025）在 Google 内部的一个应用：一个训练 Gemini 用的矩阵乘法内核，
    要按输入形状 (M, N, K) 选分块大小。以前靠专家手写的启发式（或者每个新形状都去自动调优）。
    AlphaEvolve 让 LLM 反复改写这段**启发式代码**，用真实形状上的运行时间打分，进化出来的启发式比专家的
    平均快 23%（内核），整个训练快 1%；正确性"由构造保证"——只改分块，不改算的东西。

    这里用同样的结构复现一个迷你版：
      "硬件"   = 第 13 讲 sched_lab.gemm_time 的解析模型（H100 式的参数，不是实测）；
      "形状"   = Llama-3.1-8B 和 70B（8 卡张量并行后每卡的分片）里的投影矩阵 × 各种 token 数；
      "程序"   = 一个 Python 函数 choose_config(M, N, K) → GemmConfig，写成一串 if 规则（EVOLVE-BLOCK 里的部分可以改）；
      "LLM"    = 一个**模拟的**提议器：按模板生成 SEARCH/REPLACE 形式的修改（改常数、改阈值、加规则、删规则、
                 从别的程序借规则），可选地"读"评估反馈（哪个形状离 roofline 最远）并用一点内核常识去修它；
      "评估器" = 在训练形状上算"相对专家启发式的几何平均加速比"，带两级级联（先在 6 个探针形状上粗筛）；
      "数据库" = 岛屿模型 + 按规则条数分格的精英档案（MAP-Elites 的简化）。

    **模拟的 LLM 没有任何"智能"**：它只是让我们看清 AlphaEvolve 这个循环的结构、评估器和数据库的作用、
    训练 / 测试形状的划分，以及消融实验在问什么。真正的 AlphaEvolve 用 Gemini 2.0 Flash + Pro 生成修改。

【运行】python3 evolve.py   # 自检（约 1 秒）
【依赖】numpy；../lecture13 的 sched_lab（GemmConfig、SPACE、static_valid、gemm_time、all_configs），
    它又用到 ../lecture10 的 kernel_lab.tile_order。都是只读导入。
"""

import difflib
import functools
import math
import os
import sys
from dataclasses import dataclass, field

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture13"))
sys.path.insert(0, os.path.join(HERE, "..", "lecture10"))

import sched_lab  # noqa: E402
from kernel_lab import tile_order  # noqa: E402
from sched_lab import SPACE, GemmConfig, all_configs, gemm_time, static_valid  # noqa: E402

# =============================================================================
# 0. 让第 13 讲的解析模型跑得更快（结果逐位不变）
# =============================================================================
_orig_dram = sched_lab.dram_bytes_rect


def _dram_fast(M, N, K, bm, bn, concurrent, group=None, elem=2):
    """
    sched_lab.dram_bytes_rect 的向量化版本：结果和原版**完全相同**（自检里逐个比较），只是快几十倍。

    【原版做什么】按启动顺序把块分成"波"（每波 concurrent 个块同时在跑），每波读它碰到的不同 A 块行、B 块列。
        原版用 Python 循环逐波调用 np.unique；这里把所有波排成一个二维数组（不满的最后一波补 −1），
        每行排序后数"相邻元素不同"的次数 + 1，就是每波不同值的个数（补的 −1 再减掉）。
    【Python 小知识】np.sort(a, axis=1) 对每一行排序；np.count_nonzero(np.diff(s, axis=1), axis=1) 数每行有几处"变了"。
    """
    tm, tn = math.ceil(M / bm), math.ceil(N / bn)
    pm, pn = tile_order(tm, tn, group)
    T = tm * tn
    W = max(1, concurrent)
    nw = -(-T // W)
    pad = nw * W - T

    def uniq(a):
        a2 = np.concatenate([a, np.full(pad, -1, dtype=a.dtype)]).reshape(nw, W)
        cnt = 1 + np.count_nonzero(np.diff(np.sort(a2, axis=1), axis=1), axis=1)
        if pad:
            cnt[-1] -= 1
        return cnt

    rows, cols = uniq(pm), uniq(pn)
    return int((np.minimum(rows * bm, M) + np.minimum(cols * bn, N)).sum()) * K * elem + M * N * elem


# functools.lru_cache：同样的参数第二次调用直接返回缓存的结果（"记忆化"）。gemm_time 每次都去模块里找
# dram_bytes_rect 这个名字，所以把模块里的名字换掉，gemm_time 就会用新版本。
sched_lab.dram_bytes_rect = functools.lru_cache(maxsize=None)(_dram_fast)

# =============================================================================
# 1. 形状：真实 LLM 里的矩阵乘法
# =============================================================================
# (名字) → (N, K)：输出维度、归约维度。8B：单卡；70B：8 卡张量并行后每卡的分片（第 4 讲）。
LAYERS = {
    "8B qkv": (6144, 4096), "8B o": (4096, 4096), "8B gate_up": (28672, 4096), "8B down": (4096, 14336),
    "70B/8 qkv": (1280, 8192), "70B/8 o": (8192, 1024), "70B/8 gate_up": (7168, 8192), "70B/8 down": (8192, 3584),
}
TOKENS = (1, 4, 16, 64, 128, 256, 512, 1024, 2048, 4096, 8192)   # M：decode 的批大小到 prefill 的 token 数


@dataclass(frozen=True)
class Shape:
    layer: str
    M: int
    N: int
    K: int

    def label(self):
        return f"{self.layer} M={self.M}"


def all_shapes():
    """8 种投影 × 11 种 token 数 = 88 个形状。"""
    return [Shape(name, M, N, K) for name, (N, K) in LAYERS.items() for M in TOKENS]


def split_shapes(shapes, how="random", seed=0):
    """
    划分训练 / 测试形状。

    【参数】how="random"：随机对半分（AlphaEvolve 的做法：一半形状用来进化，另一半检验泛化）；
        how="by_model"：8B 的形状训练、70B 分片的形状测试（更难：测试形状的 N、K 都没见过）
    【返回】(train, test)
    """
    if how == "by_model":
        return [s for s in shapes if s.layer.startswith("8B")], [s for s in shapes if s.layer.startswith("70B")]
    idx = np.random.default_rng(seed).permutation(len(shapes))
    half = len(shapes) // 2
    return [shapes[i] for i in sorted(idx[:half])], [shapes[i] for i in sorted(idx[half:])]


VALID = [c for c in all_configs() if static_valid(c)]


class TimeTable:
    """
    "测量"的缓存：time(cfg, shape) 第一次调用 gemm_time，之后查表。

    【为什么】进化过程中同一个 (配置, 形状) 会被问很多次；解析模型本身是确定的（没有噪声），所以可以缓存。
        真实系统里每次"测量"都要在 TPU / GPU 上跑一次内核——这正是评估成本的来源，下面用 n_measure 记下真正测了几次。
    """

    def __init__(self):
        self.t = {}
        self.n_measure = 0

    def time(self, cfg, s):
        key = (cfg, s)
        v = self.t.get(key)
        if v is None:
            v = gemm_time(cfg, s.M, s.N, s.K) if static_valid(cfg) else math.inf
            self.t[key] = v
            self.n_measure += 1
        return v

    def oracle(self, s):
        """逐形状穷举 3744 个能启动的配置，返回 (最快的配置, 时间)——"每个形状都调优"的上限。"""
        best = min(VALID, key=lambda c: self.time(c, s))
        return best, self.time(best, s)


def roofline_time(s, peak=2048 * 2 * 132 * 1.83e9, hbm=3.35e12, launch=4e-6):
    """
    roofline 下界：max(运算 / 峰值, 最少字节 / 带宽) + 启动开销（参数和 sched_lab.H100_TC 一致）。

    【用途】"离 roofline 多远"是评估器给 LLM 的反馈之一（KernelBench 也给模型看硬件规格）：比值越小的形状越值得改。
    """
    flops = 2.0 * s.M * s.N * s.K
    byts = 2.0 * (s.M * s.K + s.K * s.N + s.M * s.N)
    return max(flops / peak, byts / hbm) + launch


# =============================================================================
# 2. 程序：一串 if 规则，渲染成 Python 源码
# =============================================================================
VARS = ("M", "N", "K")
FIELDS = ("bm", "bn", "bk", "warps", "stages", "group")


@dataclass(frozen=True)
class Rule:
    """
    一条规则：conds 是若干个 "变量 <= 阈值" 的合取（全部成立才命中），cfg 是命中时返回的配置。

    【例子】Rule((("M", 16),), GemmConfig(16, 64, 64, 2, 3, 1)) 渲染成
        if M <= 16:
            return GemmConfig(bm=16, bn=64, bk=64, warps=2, stages=3, group=1)
    """
    conds: tuple
    cfg: GemmConfig


@dataclass(frozen=True)
class Program:
    """规则按顺序检查，第一条命中的规则决定配置；都不命中就用 default。"""
    rules: tuple
    default: GemmConfig

    def n_rules(self):
        return len(self.rules)


def _cfg_src(c):
    return "GemmConfig(" + ", ".join(f"{f}={v}" for f, v in zip(FIELDS, c.as_tuple())) + ")"


def render(prog):
    """
    把 Program 渲染成 Python 源码（AlphaEvolve 进化的就是这样一段代码；EVOLVE-BLOCK 标出可以改的部分）。
    """
    lines = ["def choose_config(M, N, K):", "    # EVOLVE-BLOCK-START"]
    for r in prog.rules:
        cond = " and ".join(f"{v} <= {t}" for v, t in r.conds)
        lines.append(f"    if {cond}:")
        lines.append(f"        return {_cfg_src(r.cfg)}")
    lines.append(f"    return {_cfg_src(prog.default)}")
    lines.append("    # EVOLVE-BLOCK-END")
    return "\n".join(lines) + "\n"


def static_check_program(src):
    """
    不运行程序，检查它里面**每一个**分支返回的配置能不能启动（共享内存、寄存器、warp 的块……，sched_lab.static_valid）。

    【为什么需要】评估只在训练形状上运行程序：一条从来没被训练形状走到的规则，即使配置根本不能启动，分数也看不出来，
        到了新形状上才出错——"测试没覆盖的分支"。这相当于编译器对所有分支做的检查（第 15 讲的"护栏"）。
    【做法】用 ast 找出所有 GemmConfig(...) 调用，取出常数参数，逐个调用 static_valid。
    【返回】(ok, 第一个不合法的配置或 None)
    """
    import ast
    for node in ast.walk(ast.parse(src)):
        if isinstance(node, ast.Call) and getattr(node.func, "id", None) == "GemmConfig":
            kw = {k.arg: k.value.value for k in node.keywords}
            cfg = GemmConfig(*[kw[f] for f in FIELDS])
            if not static_valid(cfg):
                return False, cfg
    return True, None


def compile_program(src):
    """执行源码，取出 choose_config（命名空间里只有 GemmConfig）。"""
    ns = {"GemmConfig": GemmConfig}
    exec(compile(src, "<evolved>", "exec"), ns)
    return ns["choose_config"]


def make_diff(old, new):
    """
    把"旧源码 → 新源码"写成 AlphaEvolve 要求 LLM 输出的 SEARCH/REPLACE 块。

    【格式】<<<<<<< SEARCH / 原来的几行 / ======= / 新的几行 / >>>>>>> REPLACE
    【做法】用 difflib 找出改动的行段，合并成一段；SEARCH 的内容必须在原文里**唯一**（否则不知道该改哪一处），
        所以从前一行开始带上下文，不唯一就继续往前、往后扩，直到唯一为止。
    【返回】块的列表（这里最多一块）；新旧相同时返回空列表
    """
    a, b = old.splitlines(keepends=True), new.splitlines(keepends=True)
    ops = [op for op in difflib.SequenceMatcher(None, a, b).get_opcodes() if op[0] != "equal"]
    if not ops:
        return []
    i1, j1 = ops[0][1], ops[0][3]
    i2, j2 = ops[-1][2], ops[-1][4]
    lo, hi = max(0, i1 - 1), i2
    while old.count("".join(a[lo:hi])) != 1 or lo == hi:
        if lo > 0:
            lo -= 1
        elif hi < len(a):
            hi += 1
        else:
            break
    search = "".join(a[lo:hi])
    replace = "".join(a[lo:i1]) + "".join(b[j1:j2]) + "".join(a[i2:hi])
    return [(search, replace)]


def apply_diff(src, blocks):
    """
    把 SEARCH/REPLACE 块应用到源码上（AlphaEvolve 系统做的事）。

    【边界】SEARCH 的内容在原文里找不到、或者出现不止一次 → 抛出 ValueError（真实系统里这样的修改直接作废）。
    """
    out = src
    for search, replace in blocks:
        n = out.count(search)
        if n != 1:
            raise ValueError("SEARCH 块在程序里找不到" if n == 0 else "SEARCH 块不唯一")
        out = out.replace(search, replace, 1)
    return out


def format_diff(blocks):
    """把 SEARCH/REPLACE 块拼成可以打印的文本。"""
    parts = []
    for s, r in blocks:
        parts.append("<<<<<<< SEARCH\n" + s + "=======\n" + r + ">>>>>>> REPLACE\n")
    return "".join(parts)


def choose(prog, s):
    """不经过源码，直接按规则选配置（和 compile_program(render(prog)) 的结果一样，自检里核对）。"""
    for r in prog.rules:
        if all(getattr(s, v) <= t for v, t in r.conds):
            return r.cfg
    return prog.default


def _implies(a, b):
    """条件 a（合取）成立时条件 b 一定成立吗？b 的每一项 (v, t) 都要被 a 里某一项 (v, t') 且 t' <= t 蕴含。"""
    return all(any(va == vb and ta <= tb for va, ta in a) for vb, tb in b)


def simplify(prog):
    """
    删掉"死规则"，不改变程序在任何输入上的行为（第 12 讲 DCE 的同一个想法）：
        1. 被前面某条规则完全挡住的规则（它的条件蕴含前面那条的条件——前面那条总是先命中）；
        2. 末尾和 default 返回同一个配置的规则（删掉后落到 default，结果一样）。
    【为什么需要】进化出来的代码常有冗余（遗传编程里叫 bloat）；AlphaEvolve 可以把"简洁"作为额外的分数，
        或者让 LLM 给可读性打分。这里用一个保证语义不变的化简，让打印出来的程序可读。
    """
    kept = []
    for r in prog.rules:
        if any(_implies(r.conds, k.conds) for k in kept):
            continue
        kept.append(r)
    while kept and kept[-1].cfg == prog.default:
        kept.pop()
    return Program(tuple(kept), prog.default)


# 专家启发式（我写的"起点"）：decode 用小块，中等批用中块，其余用 Triton 教程里常见的 128×128×64。
EXPERT = Program(rules=(Rule((("M", 32),), GemmConfig(32, 64, 64, 2, 3, 1)),
                        Rule((("M", 256),), GemmConfig(64, 128, 64, 4, 3, 8))),
                 default=GemmConfig(128, 128, 64, 4, 3, 8))


# =============================================================================
# 3. 评估器
# =============================================================================
@dataclass
class Evaluator:
    """
    给一个程序打分：在一组形状上，相对"专家启发式"的几何平均加速比（越大越好）。

    【字段】shapes：评估用的形状；table：TimeTable；probe：级联第一级的探针形状数；cascade_ratio：
        第一级的分数低于父程序的这个比例就淘汰（不再做完整评估）
    【级联】AlphaEvolve 的 evaluation cascade：先在小而便宜的测试上筛，过了再上完整的测试，省掉大部分坏程序的成本。
    【静态检查】check_all_branches=True 时先检查程序里每个分支的配置能不能启动（static_check_program）。
    【计数】n_eval_full / n_eval_probe：完整评估 / 只做了第一级的程序数；n_static_reject：被静态检查拒掉的程序数；
        bench_runs：一共"跑了多少次内核"（每个程序在每个形状上算一次；真实系统里每次都要上硬件）。
    """
    shapes: list
    table: TimeTable
    probe: int = 6
    cascade_ratio: float = 0.85
    check_all_branches: bool = True
    n_eval_full: int = 0
    n_eval_probe: int = 0
    n_static_reject: int = 0
    bench_runs: int = 0
    _base: dict = field(default_factory=dict)
    best_seen: dict = field(default_factory=dict)

    def __post_init__(self):
        for s in self.shapes:
            self._base[s] = self.table.time(choose(EXPERT, s), s)
            self.best_seen[s] = (self._base[s], choose(EXPERT, s))
        idx = np.linspace(0, len(self.shapes) - 1, self.probe).round().astype(int)
        self.probe_shapes = [self.shapes[i] for i in idx]

    def _score(self, fn, shapes):
        logs = []
        for s in shapes:
            cfg = fn(s.M, s.N, s.K)
            t = self.table.time(cfg, s)
            self.bench_runs += 1
            if not math.isfinite(t):
                return 0.0, s            # 选了一个不能启动的配置：整个程序作废（像内核编译失败）
            if t < self.best_seen[s][0]:
                self.best_seen[s] = (t, cfg)   # 记下"这个形状上目前见过的最快配置"（来自任何一个程序）
            logs.append(math.log(self._base[s] / t))
        return math.exp(float(np.mean(logs))), None

    def evaluate(self, src, parent_probe=None, cascade=True):
        """
        【返回】dict(score, probe_score, full, error)：score 是完整评估的分数（被级联淘汰时为 None）
        【边界】源码执行出错 / 返回了不能启动的配置 → score = 0，error 写明原因。
        """
        try:
            ok, bad_cfg = static_check_program(src)
            fn = compile_program(src)
        except Exception as e:  # noqa: BLE001
            return {"score": 0.0, "probe_score": 0.0, "full": False, "error": f"{type(e).__name__}"}
        if self.check_all_branches and not ok:
            self.n_static_reject += 1
            return {"score": 0.0, "probe_score": 0.0, "full": False, "error": f"有分支返回不能启动的配置 {bad_cfg.as_tuple()}"}
        ps, bad = self._score(fn, self.probe_shapes)
        if bad is not None:
            return {"score": 0.0, "probe_score": 0.0, "full": False, "error": f"不能启动：{bad.label()}"}
        if cascade and parent_probe is not None and ps < self.cascade_ratio * parent_probe:
            self.n_eval_probe += 1
            return {"score": None, "probe_score": ps, "full": False, "error": None}
        sc, bad = self._score(fn, self.shapes)
        self.n_eval_full += 1
        if bad is not None:
            return {"score": 0.0, "probe_score": ps, "full": True, "error": f"不能启动：{bad.label()}"}
        return {"score": sc, "probe_score": ps, "full": True, "error": None}

    def per_shape(self, prog):
        """
        给"LLM"看的反馈：每个训练形状一行 dict(shape, cfg, t, vs_expert, vs_roofline, vs_best, best_cfg)。

        vs_roofline = roofline 下界 / 时间（1 = 碰到了理论上限）；vs_best = 目前见过的最快时间 / 时间
        （1 = 已经和见过的最好一样快）。两个都是"越小越值得改"，但含义不同：roofline 可能根本到不了，
        "见过的最好"一定到得了——它来自评估器已经做过的测量，不需要额外的成本。
        """
        out = []
        for s in self.shapes:
            c = choose(prog, s)
            t = self.table.time(c, s)
            bt, bc = self.best_seen[s]
            out.append({"shape": s, "cfg": c, "t": t, "vs_expert": self._base[s] / t,
                        "vs_roofline": roofline_time(s) / t, "vs_best": bt / t, "best_cfg": bc})
        return out


def geomean_speedup(prog, shapes, table, base_prog=EXPERT):
    """
    程序在一组形状上相对 base_prog 的几何平均加速比（不计入评估成本，用于汇报）。

    【边界】某个形状上选了不能启动的配置（时间 = ∞）→ 返回 0.0（程序在这组形状上不可用）。
    """
    logs = []
    for s in shapes:
        t = table.time(choose(prog, s), s)
        if not math.isfinite(t):
            return 0.0
        logs.append(math.log(table.time(choose(base_prog, s), s) / t))
    return math.exp(float(np.mean(logs)))


def oracle_speedup(shapes, table, base_prog=EXPERT):
    """逐形状最优配置相对 base_prog 的几何平均加速比（上限）。"""
    logs = [math.log(table.time(choose(base_prog, s), s) / table.oracle(s)[1]) for s in shapes]
    return math.exp(float(np.mean(logs)))


def best_single_config(shapes, table):
    """所有形状共用一个配置时最好的那个（"一个尺寸走天下"的基线）。"""
    best, best_s = None, -1.0
    for c in VALID:
        sc = geomean_speedup(Program((), c), shapes, table)
        if sc > best_s:
            best, best_s = c, sc
    return Program((), best)


# =============================================================================
# 4. 模拟的 LLM：提出修改
# =============================================================================
def _neighbor(cfg, rng):
    """把配置的某一个字段改成 SPACE 里相邻的值（和第 13 讲 sched_lab.mutate 一样）。"""
    k = FIELDS[rng.integers(len(FIELDS))]
    vals = SPACE[k]
    i = vals.index(getattr(cfg, k))
    j = min(len(vals) - 1, max(0, i + (1 if rng.random() < 0.5 else -1)))
    d = dict(zip(FIELDS, cfg.as_tuple()))
    d[k] = vals[j]
    return GemmConfig(*d.values())


def _pow2(x):
    return int(2 ** round(math.log2(max(1, x))))


def _prior_config(s, rng):
    """
    "内核常识"：给一个形状猜一个合理的配置——BM 和 M 同量级（16–256），N 大就用宽的 BN，
    小块用大 BK、少 warp，大块用 8 个 warp、4 级流水线。再随机扰动 0–1 个字段（LLM 的猜测也不总是准）。
    """
    bm = min(256, max(16, _pow2(s.M)))
    bn = 128 if s.N >= 2048 else 64
    bk = 64 if bm <= 64 else 32
    warps = 8 if bm * bn >= 256 * 128 else (4 if bm * bn >= 64 * 128 else (2 if bm * bn >= 32 * 64 else 1))
    cfg = GemmConfig(bm, bn, bk, warps, 4, 1 if s.M <= 256 else 8)
    if rng.random() < 0.5:
        cfg = _neighbor(cfg, rng)
    return cfg


@dataclass
class MockLLM:
    """
    模拟的"LLM"：看父程序（和几个"灵感"程序、可选的评估反馈），提出一个修改后的新程序。

    【字段】
        knowledge：有没有"内核常识"——能不能一次提出**几个字段一起变**的连贯修改（targeted）。
                   False 时只会做单字段的小改动，相当于 AlphaEvolve 消融里的 "Small base LLM only"。
        feedback ：targeted 修改挑哪个形状下手——
                   "best_seen"：看评估反馈里"离这个形状上见过的最快配置还差多少"，差得最多的优先，并且直接借用那个配置；
                   "roofline" ：看"离 roofline 下界还差多少"，差得最多的优先；
                   None       ：不看反馈，随便挑一个形状。
        focus    ：怎样用反馈挑形状——"weighted"：按"还差多少"加权随机抽；"worst3"：只在差得最多的 3 个里挑。
        use_inspirations：是否从别的程序借规则；max_rules：规则条数上限（太长的程序难读）
    【它会做的修改】
        tweak_cfg      ：改某条规则（或 default）配置里的一个字段
        tweak_threshold：把某个条件的阈值乘 2 或除以 2
        add_rule       ：加一条单条件规则（阈值随机），配置从"原来管这片区域的规则"抄过来再改一个字段
        remove_rule    ：删一条规则
        borrow         ：从灵感程序里抄一条规则插进来
        targeted       ：（需要 knowledge）为一个形状加一条 "M <= 它的 M" 的规则，配置用见过的最快配置（如果有更快的）
                         或按内核常识猜一个（_prior_config），插在原来管这个形状的规则前面
    """
    knowledge: bool = True
    feedback: object = "best_seen"
    focus: str = "weighted"
    use_inspirations: bool = True
    max_rules: int = 12

    def propose(self, parent, inspirations, feedback, shapes, rng):
        """【返回】(新程序, 修改的种类)。feedback 是 Evaluator.per_shape 的输出（或 None）；shapes 是问题描述里的形状列表。"""
        ops = ["tweak_cfg", "tweak_threshold", "add_rule", "remove_rule"]
        weights = [0.35, 0.2, 0.2, 0.1]
        if self.use_inspirations and inspirations:
            ops.append("borrow")
            weights.append(0.15)
        if self.knowledge:
            ops.append("targeted")
            weights.append(0.35)
        w = np.array(weights) / sum(weights)
        op = ops[rng.choice(len(ops), p=w)]
        rules = list(parent.rules)
        default = parent.default
        if op == "tweak_cfg" or (op in ("tweak_threshold", "remove_rule") and not rules):
            i = rng.integers(len(rules) + 1)
            if i == len(rules):
                default = _neighbor(default, rng)
            else:
                rules[i] = Rule(rules[i].conds, _neighbor(rules[i].cfg, rng))
            op = "tweak_cfg"
        elif op == "tweak_threshold":
            i = rng.integers(len(rules))
            conds = list(rules[i].conds)
            j = rng.integers(len(conds))
            v, t = conds[j]
            conds[j] = (v, max(1, t * 2 if rng.random() < 0.5 else t // 2))
            rules[i] = Rule(tuple(conds), rules[i].cfg)
        elif op == "remove_rule":
            del rules[rng.integers(len(rules))]
        elif op == "add_rule":
            v = VARS[rng.choice(3, p=[0.6, 0.2, 0.2])]
            t = int(2 ** rng.integers(0, 15))
            pos = int(rng.integers(len(rules) + 1))
            probe = Shape("", t if v == "M" else 1, t if v == "N" else 4096, t if v == "K" else 4096)
            base = choose(Program(tuple(rules[pos:]), default), probe)
            rules.insert(pos, Rule(((v, t),), _neighbor(base, rng)))
        elif op == "borrow":
            donor = inspirations[rng.integers(len(inspirations))]
            if donor.rules:
                r = donor.rules[rng.integers(len(donor.rules))]
                rules.insert(int(rng.integers(len(rules) + 1)), r)
        elif op == "targeted":
            fb = None
            if self.feedback is not None and feedback:
                key = "vs_best" if self.feedback == "best_seen" else "vs_roofline"
                if self.focus == "worst3":
                    worst = sorted(feedback, key=lambda x: x[key])[:3]    # 只盯"最该改"的 3 个形状
                    fb = worst[rng.integers(len(worst))]
                else:                                                     # 按"还差多少"加权抽一个形状
                    wts = np.array([1.0 / x[key] - 1.0 + 0.05 for x in feedback])
                    fb = feedback[rng.choice(len(feedback), p=wts / wts.sum())]
                s = fb["shape"]
            else:
                s = shapes[rng.integers(len(shapes))]
            conds = [("M", s.M)]
            if rng.random() < 0.3:
                conds.append(("N", _pow2(s.N) if _pow2(s.N) >= s.N else 2 * _pow2(s.N)))
            pos = len(rules)
            for k, r in enumerate(rules):                             # 插在"原来管它的那条规则"前面
                if all(getattr(s, vv) <= tt for vv, tt in r.conds):
                    pos = k
                    break
            if fb is not None and self.feedback == "best_seen" and fb["vs_best"] < 1.0:
                cfg = fb["best_cfg"]
            else:
                cfg = _prior_config(s, rng)
            rules.insert(pos, Rule(tuple(conds), cfg))
        if len(rules) > self.max_rules:
            del rules[rng.integers(len(rules))]
        return Program(tuple(rules), default), op


# =============================================================================
# 5. 程序数据库：岛屿 + 精英档案
# =============================================================================
@dataclass
class Database:
    """
    【岛屿】n_islands 个互相隔离的种群（FunSearch / AlphaEvolve 的岛屿模型）：每个子程序回到父程序所在的岛；
        每隔 migrate_every 次，把每个岛最好的程序复制到下一个岛。隔离保持多样性，迁移让好的想法传播。
    【精英档案】按"规则条数"分格，每格只留分数最高的程序（MAP-Elites 的一维简化）；"灵感"从档案里抽。
    【选父程序】在一个岛里做 3 选 1 的锦标赛（偏向高分，但不总是选最好的——保留探索）。
    """
    n_islands: int = 4
    island_cap: int = 20
    migrate_every: int = 50
    islands: list = field(default_factory=list)
    elites: dict = field(default_factory=dict)
    added: int = 0

    def seed(self, prog, res):
        self.islands = [[(prog, res)] for _ in range(self.n_islands)]
        self.elites = {prog.n_rules(): (prog, res)}

    def sample(self, rng):
        isl = int(rng.integers(self.n_islands))
        pop = self.islands[isl]
        picks = [pop[rng.integers(len(pop))] for _ in range(min(3, len(pop)))]
        parent = max(picks, key=lambda x: x[1]["score"])
        insp = [p for p, _ in self.elites.values() if p is not parent[0]]
        if len(insp) > 2:
            insp = [insp[i] for i in rng.choice(len(insp), 2, replace=False)]
        return isl, parent, insp

    def add(self, isl, prog, res):
        pop = self.islands[isl]
        pop.append((prog, res))
        pop.sort(key=lambda x: -x[1]["score"])
        del pop[self.island_cap:]
        k = prog.n_rules()
        if k not in self.elites or res["score"] > self.elites[k][1]["score"]:
            self.elites[k] = (prog, res)
        self.added += 1
        if self.n_islands > 1 and self.added % self.migrate_every == 0:
            bests = [pop[0] for pop in self.islands]
            for i, b in enumerate(bests):
                self.islands[(i + 1) % self.n_islands].append(b)

    def best(self):
        return max((pop[0] for pop in self.islands), key=lambda x: x[1]["score"])


# =============================================================================
# 6. 进化循环
# =============================================================================
def evolve(evaluator, budget, seed=0, llm=None, evolution=True, n_islands=4, cascade=True, initial=EXPERT,
           log_every=None):
    """
    AlphaEvolve 的主循环（单线程版）：
        parent, inspirations = database.sample()
        diff = llm.generate(prompt(parent, inspirations, feedback))
        child = apply_diff(parent, diff)
        results = evaluator.execute(child)
        database.add(child, results)

    【参数】budget：最多提出多少个程序；evolution=False 时每次都从**初始程序**出发改（"No evolution"消融）；
        n_islands：岛屿数；cascade：是否用两级级联
    【返回】dict：best（最好的 (程序, 结果)；结果里有产生它的 diff 和修改种类）、history（第 i 个程序之后的最好分数）、
        accepted / tried（每种修改"比父程序好" / 被提出的次数）、improvements（全局最好被刷新的记录）、
        n_added / n_added_bad_branch（进数据库的程序数 / 其中带不合法分支的个数）、db
    """
    rng = np.random.default_rng(seed)
    llm = llm or MockLLM()
    db = Database(n_islands=n_islands)
    init_res = evaluator.evaluate(render(initial))
    db.seed(initial, init_res)
    hist, accepted, tried = [], {}, {}
    improvements = []            # 全局最好分数每次被刷新时：(第几个程序, 新分数, 修改的种类)
    n_added = n_added_bad = 0    # 进了数据库的程序数；其中"有分支返回不能启动的配置"的个数（只在关掉分支检查时可能非 0）
    for it in range(budget):
        if evolution:
            isl, (parent, pres), insp = db.sample(rng)
        else:
            isl, parent, pres, insp = 0, initial, init_res, []
        fb = evaluator.per_shape(parent) if llm.feedback is not None else None
        child, op = llm.propose(parent, insp, fb, evaluator.shapes, rng)
        tried[op] = tried.get(op, 0) + 1
        psrc, csrc = render(parent), render(child)
        blocks = make_diff(psrc, csrc)
        if not blocks:                                   # 修改没有改变程序
            hist.append(db.best()[1]["score"])
            continue
        src = apply_diff(psrc, blocks)                   # 系统把 diff 应用到父程序上
        res = evaluator.evaluate(src, parent_probe=pres["probe_score"], cascade=cascade)
        if res["score"] is not None and res["score"] > 0:
            res["op"], res["diff"] = op, blocks
            prev_best = db.best()[1]["score"]
            if res["score"] > pres["score"]:
                accepted[op] = accepted.get(op, 0) + 1
            db.add(isl if evolution else int(rng.integers(db.n_islands)), child, res)
            n_added += 1
            n_added_bad += not static_check_program(src)[0]
            if res["score"] > prev_best:
                improvements.append((it + 1, res["score"], op))
        hist.append(db.best()[1]["score"])
        if log_every and (it + 1) % log_every == 0:
            print(f"    第 {it + 1:4d} 个程序：最好分数 {hist[-1]:.3f}")
    return {"best": db.best(), "history": hist, "accepted": accepted, "tried": tried, "db": db,
            "improvements": improvements, "n_added": n_added, "n_added_bad_branch": n_added_bad}


# =============================================================================
# 自检
# =============================================================================
def _selftest():
    rng = np.random.default_rng(0)
    for _ in range(200):          # 快版 DRAM 模型和原版逐位一致
        M = int(rng.choice([1, 3, 16, 100, 512, 2048])); N = int(rng.choice([64, 1000, 4096, 6144]))
        bm = int(rng.choice(SPACE["bm"])); bn = int(rng.choice(SPACE["bn"]))
        c = int(rng.choice([1, 7, 132, 264, 1056])); g = [None, 4, 8][rng.integers(3)]
        assert _orig_dram(M, N, 4096, bm, bn, c, g) == _dram_fast(M, N, 4096, bm, bn, c, g)
    shapes = all_shapes()
    assert len(shapes) == 88
    table = TimeTable()
    # 渲染 → 编译 → 选配置，和直接按规则选一致；diff 可以还原
    src = render(EXPERT)
    fn = compile_program(src)
    for s in shapes:
        assert fn(s.M, s.N, s.K) == choose(EXPERT, s)
    llm = MockLLM()
    p = EXPERT
    for _ in range(300):
        q, _ = llm.propose(p, [EXPERT], None, shapes, rng)
        assert apply_diff(render(p), make_diff(render(p), render(q))) == render(q)
        p = q
    # 一个短的进化：分数不低于起点，且最好的程序在训练形状上确实更快
    train, _ = split_shapes(shapes)
    ev = Evaluator(train, table)
    out = evolve(ev, 150, seed=0)
    best_prog, best_res = out["best"]
    assert best_res["score"] >= 1.0
    assert abs(geomean_speedup(best_prog, train, table) - best_res["score"]) < 1e-9
    assert oracle_speedup(train[:5], table) >= geomean_speedup(best_prog, train[:5], table) - 1e-12
    print(f"evolve 自检通过（150 个程序后训练集分数 {best_res['score']:.3f}）。")


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