"""
=============================================================================
第 7 讲 · 01：用"演化 + 自动评估器"发现算法 —— FunSearch / AlphaEvolve 的缩小版
=============================================================================

对应讲义：lectures/lecture07_open_ended_evolution.md 第 4 节

-----------------------------------------------------------------------------
【问题：在线装箱（online bin packing）】
-----------------------------------------------------------------------------
物品一个接一个地到来（大小按 Weibull 分布，集中在 40 左右），箱子容量 100。每来一个物品，必须马上决定放进哪个箱子
（放得下的箱子里选一个；都放不下就开一个新箱子），之后不能再挪动。目标：用的箱子越少越好。
经典的启发式规则：
    First Fit：放进第一个放得下的箱子；
    Best Fit ：放进放下之后剩余空间最小的箱子（"塞得最紧"）。
FunSearch（Romera-Paredes et al. 2023）就在这个问题上，用 LLM 演化出了比 Best Fit 更好的规则。

-----------------------------------------------------------------------------
【怎样"演化"一个算法】
-----------------------------------------------------------------------------
    · 程序：一个"优先级函数" priority(s, r)：s 是物品大小，r 是某个箱子的剩余空间（程序里还可以用
      g = r − s，放进去之后剩下的空隙）。每来一个物品，放进优先级最高的那个（放得下的）箱子。
      例如 Best Fit 就是 priority = −(r − s)。程序可以用 + − ×、受保护的除法 pdiv、min、max，
      以及条件运算 ifgt(a, b, c, d)（a > b 时取 c，否则取 d）。
      程序写成一段 Python 表达式（字符串），用 eval 编译成函数 —— **程序就是代码**，这正是 FunSearch / AlphaEvolve 的做法。
    · 评估器（evaluator）：在一批随机实例上模拟装箱，分数 = −（用的箱子数比下界多出的百分比）。
      下界 = ⌈所有物品大小之和 / 容量⌉。分数越接近 0 越好。评估是全自动的，这是这类方法成立的前提。
    · 提议新程序：真实系统里是 LLM 读入几个已有的好程序、写出一个改进的版本。这里没有 LLM，
      用几种随机操作代替：**局部修改**（在原程序上加一个小的修正项，像 LLM 对代码做的局部编辑）、
      **变异**（随机替换表达式的一部分、改一个常数）和**交叉**（把两个好程序的部分拼起来）。
      LLM 提议的修改比随机变异"聪明"得多（它懂语义、能借鉴文献），所以真实系统需要的评估次数少得多。
    · 程序数据库：保存所有评估过的程序。三种搜索方式：
          随机搜索       ：每次随机写一个新程序（不利用已有的结果）；
          爬山（greedy） ：只保留当前最好的一个程序，每次对它变异，更好就替换；
          岛屿演化       ：4 个"岛"各自保留一群好程序，每次从某个岛里挑两个好程序交叉 / 变异；
                          定期把最差的岛清空、用最好的岛的程序重新播种（FunSearch 的做法，保持多样性）。

-----------------------------------------------------------------------------
【两个演示】
-----------------------------------------------------------------------------
演示 1  三种搜索方式在同样的评估预算下找到的最好程序，和 First Fit / Best Fit 比较（在**没见过的**实例上）。
演示 2  评估器的设计：如果评估器只用 2 个固定的小实例（便宜），演化出的程序会"过拟合"评估器吗？

运行时间：约 6 秒（只用标准库）。
"""

import math
import random

CAPACITY = 100
WEIBULL_SCALE, WEIBULL_SHAPE = 45.0, 3.0   # FunSearch 论文用的 Weibull 分布
VARS = ("s", "r", "g")                 # 物品大小、箱子剩余空间、放进去之后的空隙 g = r − s
BINARY = ("+", "-", "*", "pdiv", "min", "max")
CONSTS = (1.0, 2.0, 5.0, 10.0, 20.0, 50.0, 100.0)


# =============================================================================
# 第一部分：装箱模拟与评估器
# =============================================================================
def weibull_size(rng):
    """
    按 Weibull(尺度 45, 形状 3) 分布抽一个物品大小，取整并截断到 [1, 100]。

    【思路】FunSearch 在 Weibull 分布的实例上报告了对 Best Fit 的大幅改进。Weibull 的大小集中在 40 左右，
        两个物品常常放不满一个箱子、三个又放不下 —— "怎样留空间"变得很关键，这正是好的启发式能发挥作用的地方。
    【Python 小知识】逆变换采样：u 是 (0, 1) 上的均匀随机数，尺度 × (−log(1 − u))^(1/形状) 就服从 Weibull 分布。
        这里自己写，而不用 random.weibullvariate，是为了保证不同 Python 版本的结果完全相同。
    """
    u = rng.random()
    x = WEIBULL_SCALE * (-math.log(1.0 - u)) ** (1.0 / WEIBULL_SHAPE)
    return min(CAPACITY, max(1, int(round(x))))


def make_instances(rng, n_instances, n_items):
    """随机生成 n_instances 个实例，每个是 n_items 个物品大小（Weibull 分布）。"""
    return [[weibull_size(rng) for _ in range(n_items)] for _ in range(n_instances)]


def pack(items, priority):
    """
    用优先级函数在线装箱，返回用了几个箱子。

    【思路】remaining 记录每个已开的箱子的剩余空间。每来一个物品 s：
        在所有放得下的箱子（r ≥ s）里，选 priority(s, r) 最大的那个；并列时选编号小的。
        都放不下就开一个新箱子。
    【参数】items：物品大小的列表；priority：函数 (s, r) → 数值
    【返回】int：箱子数
    【边界】priority 可能返回 nan 或报错（例如演化出的奇怪表达式）。这里把出错的程序当作"很差的程序"：
        评估函数会捕获异常，给它最低分。
    """
    remaining = []
    for s in items:
        best_i, best_p = -1, -math.inf
        for i, r in enumerate(remaining):
            if r >= s:
                p = priority(s, r)
                if p > best_p:
                    best_i, best_p = i, p
        if best_i < 0:
            remaining.append(CAPACITY - s)
        else:
            remaining[best_i] -= s
    return len(remaining)


def evaluate(priority, instances):
    """
    评估器：在一批实例上，用的箱子数比下界多出的平均百分比，取负数作为分数（越接近 0 越好）。

    【返回】float（例如 −4.2 表示平均比下界多用 4.2% 的箱子）；程序出错时返回 −100。
    【边界】下界 ⌈Σ大小 / 容量⌉ 不一定能达到，所以即使是最优解，分数也通常是负的。
    """
    try:
        excess = []
        for items in instances:
            lb = math.ceil(sum(items) / CAPACITY)
            excess.append(100.0 * (pack(items, priority) - lb) / lb)
        return -math.fsum(excess) / len(excess)
    except (ZeroDivisionError, OverflowError, ValueError, TypeError):
        return -100.0


# =============================================================================
# 第二部分：程序 = 表达式树 ↔ Python 代码
# =============================================================================
def pdiv(a, b):
    """受保护的除法：分母接近 0 时返回分子（演化出的程序里随时可能出现除以 0）。"""
    return a / b if abs(b) > 1e-9 else a


def ifgt(a, b, c, d):
    """条件运算：a > b 时返回 c，否则返回 d。有了它，程序才能表达"分情况处理"的规则。"""
    return c if a > b else d


def to_code(tree):
    """
    把表达式树转成 Python 代码字符串。

    【表示】树是嵌套的 tuple：("var", "s")、("const", 2.0)、(运算符, 子树, 子树)、("ifgt", 子树 × 4)。
        例如 Best Fit 的 −(r − s) 写成 ("-", ("const", 0.0), ("-", ("var", "r"), ("var", "s")))。
        变量 g 写成代码时展开成 (r - s)。
    """
    kind = tree[0]
    if kind == "var":
        return "(r - s)" if tree[1] == "g" else tree[1]
    if kind == "const":
        return repr(tree[1])
    args = [to_code(child) for child in tree[1:]]
    if kind in ("pdiv", "min", "max", "ifgt"):
        return f"{kind}({', '.join(args)})"
    return f"({args[0]} {kind} {args[1]})"


def compile_tree(tree):
    """
    把表达式树编译成真正的 Python 函数。

    【Python 小知识】eval("lambda s, r: ...", 命名空间) 把一段字符串当作 Python 代码执行，得到一个函数。
        eval 很强大，也很危险：**绝不能对不可信的输入使用 eval**（别人可以借此执行任意代码）。
        这里的字符串全是我们自己的程序生成的，只包含 s、r、数字、+ − ×、pdiv、min、max、ifgt，所以是安全的。
        AlphaEvolve 这类系统执行 LLM 写的代码时，必须放在沙盒里运行（讲义第 5 节）。
    """
    return eval(f"lambda s, r: {to_code(tree)}", {"pdiv": pdiv, "min": min, "max": max, "ifgt": ifgt})


def random_tree(rng, depth):
    """随机生成一棵深度不超过 depth 的表达式树（25% 的内部节点是 ifgt）。"""
    if depth == 0 or rng.random() < 0.3:
        if rng.random() < 0.6:
            return ("var", VARS[rng.randrange(len(VARS))])
        return ("const", CONSTS[rng.randrange(len(CONSTS))])
    if rng.random() < 0.25:
        return ("ifgt",) + tuple(random_tree(rng, depth - 1) for _ in range(4))
    return (BINARY[rng.randrange(len(BINARY))], random_tree(rng, depth - 1), random_tree(rng, depth - 1))


def subtrees(tree, path=()):
    """
    列出树里所有子树的位置（从根出发的路径，例如 (1, 2) 表示"第 1 个孩子的第 2 个孩子"）。

    【Python 小知识】yield 让函数变成"生成器"：每 yield 一次交出一个值，调用者可以用 for 或 list() 逐个取出。
        yield from 把另一个生成器的值原样交出（这里用来递归地遍历子树）。
    """
    yield path
    if tree[0] not in ("var", "const"):
        for i in range(1, len(tree)):
            yield from subtrees(tree[i], path + (i,))


def get_at(tree, path):
    for i in path:
        tree = tree[i]
    return tree


def replace_at(tree, path, new):
    """返回一棵新树：把 path 位置的子树换成 new（原来的树不变）。"""
    if not path:
        return new
    parts = list(tree)
    parts[path[0]] = replace_at(tree[path[0]], path[1:], new)
    return tuple(parts)


def size(tree):
    return 1 if tree[0] in ("var", "const") else 1 + sum(size(child) for child in tree[1:])


MAX_SIZE = 30


def mutate(rng, tree):
    """
    随机修改一个程序（代替"LLM 提议的修改"）。

    【四种修改】
        · 40%：局部修改 —— 在原程序上加一个小的修正项：原程序 + c × (一个小表达式)。
          这样的修改不会破坏原程序的大部分行为，最像 LLM 对代码做的"小改动"；
        · 其余：换掉一棵随机子树、把一个常数乘以 0.5~2、换掉一个二元运算符。
    【边界】树太大（超过 MAX_SIZE 个节点）时，改为换成一棵小树，防止程序无限膨胀（"bloat"）。
    """
    if rng.random() < 0.4:
        sign = 1 if rng.random() < 0.5 else -1
        out = ("+", tree, ("*", ("const", sign * CONSTS[rng.randrange(len(CONSTS))]), random_tree(rng, 2)))
        return out if size(out) <= MAX_SIZE else random_tree(rng, 3)
    paths = list(subtrees(tree))
    path = paths[rng.randrange(len(paths))]
    node = get_at(tree, path)
    u = rng.random()
    if node[0] == "const" and u < 0.5:
        new = ("const", round(node[1] * (0.5 + 1.5 * rng.random()), 2))
    elif node[0] in BINARY and u < 0.5:
        new = (BINARY[rng.randrange(len(BINARY))], node[1], node[2])
    else:
        new = random_tree(rng, 2)
    out = replace_at(tree, path, new)
    return out if size(out) <= MAX_SIZE else random_tree(rng, 3)


def crossover(rng, a, b):
    """把程序 b 的一棵随机子树，接到程序 a 的一个随机位置上（代替"LLM 综合两个好程序"）。"""
    pa = list(subtrees(a))
    pb = list(subtrees(b))
    out = replace_at(a, pa[rng.randrange(len(pa))], get_at(b, pb[rng.randrange(len(pb))]))
    return out if size(out) <= MAX_SIZE else mutate(rng, a)


BEST_FIT = ("-", ("const", 0.0), ("-", ("var", "r"), ("var", "s")))    # priority = −(r − s)
BUDGET = 1500


# =============================================================================
# 第三部分：三种搜索方式
# =============================================================================
def search(method, rng, train, budget):
    """
    在评估预算 budget（评估多少个程序）内搜索，返回 (最好的树, 它在训练实例上的分数, 每 100 次评估时的最好分数)。

    【参数】method："random" / "greedy" / "islands"；train：评估器用的实例
    【思路】三种方法都从 Best Fit 出发（真实系统也常从一个已知的基线程序开始演化）。
    """
    seed_score = evaluate(compile_tree(BEST_FIT), train)
    best_tree, best_score = BEST_FIT, seed_score
    history = []
    if method == "islands":
        islands = [[(seed_score, BEST_FIT)] for _ in range(4)]
    current = (seed_score, BEST_FIT)
    for t in range(1, budget + 1):
        if method == "random":
            cand = random_tree(rng, 3)
        elif method == "greedy":
            cand = mutate(rng, current[1])
        else:
            isl = islands[rng.randrange(4)]
            parents = sorted(isl, key=lambda x: -x[0])[:5]          # 岛上最好的 5 个程序
            a = parents[rng.randrange(len(parents))][1]
            b = parents[rng.randrange(len(parents))][1]
            cand = crossover(rng, a, b) if rng.random() < 0.5 else mutate(rng, a)
        score = evaluate(compile_tree(cand), train)
        if method == "greedy" and score >= current[0]:
            current = (score, cand)
        if method == "islands":
            isl.append((score, cand))
            isl.sort(key=lambda x: -x[0])
            del isl[10:]                                            # 每个岛只保留最好的 10 个
            if t % 150 == 0:                                        # 定期清空最差的岛，用最好的岛重新播种
                order = sorted(range(4), key=lambda i: islands[i][0][0])
                islands[order[0]] = [islands[order[-1]][0]]
        if score > best_score:
            best_tree, best_score = cand, score
        if t % 100 == 0:
            history.append(best_score)
    return best_tree, best_score, history


def main():
    rng = random.Random(170)
    train = make_instances(rng, 4, 250)         # 评估器用的实例（演化时看得到）
    test = make_instances(rng, 20, 1000)        # 没见过的实例（只用来报告）
    first_fit = lambda s, r: 0.0                 # 所有箱子优先级相同 → 并列时选编号最小的 = First Fit
    best_fit = compile_tree(BEST_FIT)
    print("=" * 78)
    print(f"演示 1：在线装箱的启发式规则 —— 三种搜索方式，各评估 {BUDGET} 个程序")
    print("=" * 78)
    print(f"基线（测试集 20 个实例 × 1000 个物品）：First Fit {evaluate(first_fit, test):.3f}，"
          f"Best Fit {evaluate(best_fit, test):.3f}（分数 = −比下界多用的箱子百分比，越接近 0 越好）\n")
    print(pad_row(["搜索方式", "训练集上最好分数", "测试集分数", "每 300 次评估时的最好训练分数"]))
    results = {}
    for method, label in (("random", "随机搜索"), ("greedy", "爬山"), ("islands", "岛屿演化")):
        tree, score, hist = search(method, random.Random(171), train, BUDGET)
        results[method] = tree
        print(pad_row([label, f"{score:.3f}", f"{evaluate(compile_tree(tree), test):.3f}",
                       " → ".join(f"{h:.2f}" for h in hist[2::3])]))
    print("\n岛屿演化找到的程序：priority(s, r) =", to_code(results["islands"]))
    print("\n观察：")
    print("  · 随机搜索几乎没有改进：不利用已有的好程序，就是在一个巨大的空间里瞎找。")
    print("  · 爬山和岛屿演化都找到了在**没见过的**实例上胜过 Best Fit 的规则（多用的箱子从 4.63% 降到 4.3% 和 4.2% 左右）。")
    print("    改进来自\"对空隙 g 分情况处理\"的条件规则 —— 这是人很难凭直觉设计好的（编写时手写过一条\"避免留下小空隙\"的")
    print("    规则，反而比 Best Fit 差）。")
    print("  · 岛屿演化的**训练**分数比爬山差，**测试**分数却更好：评估器只用了 4 个实例，本身有噪声，")
    print("    爬山会把这几个实例的\"运气\"也拟合进去。更大、更多样的评估集给出更可靠的信号。")
    print("  · 演化出的程序很难读（表达式不断膨胀）。FunSearch 偏好短程序，AlphaEvolve 强调结果要可解释、可检验。")
    print("  · 改进幅度远小于 FunSearch（4.0% → 0.7%）：真实的 LLM 提议的修改比这里的随机修改聪明得多，")
    print("    而且 FunSearch 评估了上百万个程序。这里只评估了 1500 个。")

    print("\n" + "=" * 78)
    print("演示 2：评估器的设计 —— 只用 2 个固定的小实例评估，会发生什么？")
    print("=" * 78)
    tiny = make_instances(random.Random(172), 2, 40)
    tree, score, _ = search("greedy", random.Random(173), tiny, BUDGET)
    print(f"  用 2 个小实例（各 40 个物品）当评估器：训练分数 {score:.3f}，测试集分数 {evaluate(compile_tree(tree), test):.3f}")
    print(f"  对照：Best Fit 在这 2 个小实例上的分数 {evaluate(best_fit, tiny):.3f}，测试集 {evaluate(best_fit, test):.3f}")
    print("  演化出的程序：priority(s, r) =", to_code(tree))
    print("\n观察：")
    print("  · 化简一下这个程序：ifgt(r, r, ...) 永远取后一项 2.0；可行的箱子总有 s ≤ r，所以 ifgt(s, r, ...) 取 100。")
    print("    于是优先级恒为 100，所有箱子一样 —— 这就是 First Fit！")
    print("  · 在那 2 个小实例上，First Fit 碰巧比 Best Fit 好，于是演化\"发现\"了它；但在测试集上，")
    print("    它的表现就是 First Fit 的水平，比 Best Fit 差得多。")
    print("  · **评估器就是目标**。评估器太小、太窄，搜索就会找到\"在评估器上好\"而不是\"真的好\"的程序 ——")
    print("    和第 3 讲 verifier 被攻破、第 6 讲 reward hacking 是同一个问题。AlphaEvolve 用多个评估器、")
    print("    逐级加难的评估级联（evaluation cascade）来减少这种风险。\n")


def pad_row(cells, widths=(10, 18, 12, 40)):
    """简单的表格行：每列左对齐到固定的显示宽度。"""
    import unicodedata
    out = ""
    for c, w in zip(cells, widths):
        width = sum(2 if unicodedata.east_asian_width(ch) in ("W", "F") else 1 for ch in c)
        out += c + " " * max(1, w - width)
    return out


if __name__ == "__main__":
    main()