"""
=============================================================================
第 7 讲 · 02：Darwin Gödel Machine —— 开放式的存档 + 自己修改自己
=============================================================================

对应讲义：lectures/lecture07_open_ended_evolution.md 第 3 节（Zhang et al. 2025）

-----------------------------------------------------------------------------
【一个风格化的模拟】
-----------------------------------------------------------------------------
一个 coding agent 的"设计"= 它装了哪些改进（10 个可选的改进，名字取自 DGM 论文里发现的那类改进）。
它在一个 benchmark 上的真实得分由设计决定（见 TRUE_SCORE 的说明）。关键的结构是**垫脚石**（stepping stones）：
    · 有些改进单独加上会让得分**下降**：例如"多次尝试"单独加上只是更容易超时；
    · 但和别的改进组合起来很有用："多次尝试 + 用另一个模型挑选最好的尝试"。
    · 只接受"立刻变好"的修改（贪心），就永远到不了那些组合。

DGM 的循环（每一轮）：
    1. 从存档里选一个父 agent：概率 ∝ 它的（评估）得分 / (1 + 它已有的子 agent 数) —— 好的优先，没怎么被探索过的也优先；
    2. 父 agent **修改自己**：读自己在 benchmark 上的失败记录，提议一个改进（加上或去掉一个功能），
       修改可能把 agent 改坏（"代码写崩了"，得分掉到很低）；
    3. 在 benchmark 上评估新 agent（评估有噪声）；
    4. 能正常运行的新 agent 都放进存档（即使它比父 agent 差）。
修改的质量取决于"修改者"自己有多强：越强的 coding agent，越会诊断失败、写出对的修改、越少把自己改坏。
这就是"自我改进"的复利：agent 变强 → 修改自己的能力也变强。

三个版本（论文里的消融）：
    DGM                       ：有存档；修改者 = 父 agent 自己。
    DGM w/o open-ended        ：没有存档，只保留"当前版本"：新版本更好才替换（贪心爬山）。
    DGM w/o self-improve      ：有存档；但修改者永远是最初的那个 agent（修改能力不随进步而提高）。

【重要】得分函数、修改质量和"能力"的关系都是**设定**的，用来说明 DGM 的两个消融为什么会有那样的结果，
不是对真实系统的定量模拟。

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

import math
import random

FEATURES = ("按行查看文件", "字符串替换编辑", "多次尝试", "另一个模型挑选", "上下文摘要",
            "自动生成测试", "记住失败的尝试", "并行运行测试", "补丁格式检查", "长输出截断")
BASE = 0.20
SOLO = (0.04, 0.03, -0.08, -0.06, 0.03, 0.01, 0.01, -0.04, 0.02, 0.02)      # 每个改进单独的效果
SYNERGY = {(1, 0): 0.05,            # 字符串替换编辑 + 按行查看：编辑更精确
           (2, 3): 0.28,            # 多次尝试 + 另一个模型挑选：最大的提升，但两者单独都是负的
           (2, 5): 0.04,            # 多次尝试 + 自动生成测试：用测试来挑选
           (5, 7): 0.13,            # 自动生成测试 + 并行运行测试（并行运行单独是负的）
           (4, 6): 0.03}            # 上下文摘要 + 记住失败的尝试


def true_score(design):
    """
    一个设计的真实得分（0~1，像 SWE-bench 的解决率）。

    【参数】design：frozenset，装了哪些改进（下标）
    【思路】得分 = 基础 0.20 + 各改进单独的效果 + 组合的额外效果。几个单独为负的改进
        （多次尝试、另一个模型挑选、并行运行测试）只有和搭档一起才有正的总效果 —— 它们是"垫脚石"。
    """
    s = BASE + math.fsum(SOLO[i] for i in design)
    s += math.fsum(v for (a, b), v in SYNERGY.items() if a in design and b in design)
    return min(1.0, max(0.0, s))


OPTIMUM = max(true_score(frozenset(i for i in range(10) if mask >> i & 1)) for mask in range(1 << 10))


def propose(modifier_score, parent, rng):
    """
    "修改者"提议对父设计的一个修改：加上或去掉一个改进。

    【思路】修改者读父 agent 的失败记录，诊断问题：
        · 以概率 skill 提出"立刻就能看到好处"的修改（在所有单步修改里选真实增益最大的那个）；
        · 否则提出一个随机的修改（加上或去掉一个随机的改进）—— 这是探索，也可能是诊断错了。
        · 修改有概率 break_prob 把 agent 改坏（得分掉到 0.02）。
        skill 和 break_prob 都由修改者自己的得分决定：修改者越强，诊断越准、越少改坏。
    【参数】modifier_score：修改者的（真实）得分；parent：父设计；rng
    【返回】(新设计, 是否改坏了)
    """
    skill = min(0.8, 0.2 + 1.2 * (modifier_score - BASE))
    break_prob = max(0.05, 0.35 - 1.0 * (modifier_score - BASE))
    if rng.random() < skill:
        options = [parent ^ {i} for i in range(10)]          # 【Python 小知识】集合的 ^：对称差，有就去掉、没有就加上
        new = max(options, key=lambda d: (true_score(d), -sorted(d ^ parent)[0]))
    else:
        new = parent ^ {rng.randrange(10)}
    return frozenset(new), rng.random() < break_prob


def run(variant, rng, iterations=150, noise=0.02):
    """
    跑一次 DGM 的某个版本，返回每一轮之后"存档里真实得分最高的设计"的得分。

    【参数】variant："dgm" / "no_open_ended" / "no_self_improve"；noise：评估噪声的标准差
    """
    base = frozenset()
    evaluate = lambda d, broken: (0.02 if broken else true_score(d)) + rng.gauss(0, noise)
    archive = [{"design": base, "eval": evaluate(base, False), "broken": False, "children": 0}]
    current = archive[0]
    best_curve = []
    for _ in range(iterations):
        if variant == "no_open_ended":
            parent = current
        else:
            weights = [max(0.01, a["eval"]) / (1 + a["children"]) for a in archive]
            u, acc = rng.random() * math.fsum(weights), 0.0
            for a, w in zip(archive, weights):
                acc += w
                if u < acc:
                    parent = a
                    break
            else:
                parent = archive[-1]
        parent["children"] += 1
        modifier = base if variant == "no_self_improve" else parent["design"]
        design, broken = propose(true_score(modifier), parent["design"], rng)
        child = {"design": design, "eval": evaluate(design, broken), "broken": broken, "children": 0}
        if variant == "no_open_ended":
            if not broken and child["eval"] >= current["eval"]:
                current = child
                archive.append(child)
        elif not broken:
            archive.append(child)                               # 能运行的都进存档，即使比父 agent 差
        best_curve.append(max(true_score(a["design"]) for a in archive if not a["broken"]))
    return best_curve, archive


def main():
    print("=" * 78)
    print(f"DGM 的消融（最好的设计的真实得分 = {OPTIMUM:.2f}，初始 agent = {BASE:.2f}；每个版本跑 30 次取平均）")
    print("=" * 78)
    variants = (("DGM（存档 + 自己修改自己）", "dgm"), ("DGM w/o open-ended（只保留最新）", "no_open_ended"),
                ("DGM w/o self-improve（修改者不变）", "no_self_improve"))
    checkpoints = (20, 40, 80, 150)
    print(f"{'':34}" + "".join(f"第 {c} 轮".rjust(10) for c in checkpoints) + "   用到\"多次尝试 + 挑选\"的比例")
    for name, key in variants:
        curves, combo = [], 0
        for seed in range(30):
            curve, archive = run(key, random.Random(180 + seed))
            curves.append(curve)
            best = max((a for a in archive if not a["broken"]), key=lambda a: true_score(a["design"]))
            combo += {2, 3} <= best["design"]
        means = [math.fsum(c[t - 1] for c in curves) / len(curves) for t in checkpoints]
        width = sum(2 if ord(ch) > 127 else 1 for ch in name)
        print(name + " " * max(1, 34 - width) + "".join(f"{m:10.3f}" for m in means) + f"   {combo / 30:.0%}")

    print("\n观察：")
    print("  · 只保留最新版本（贪心）前期最快：每一轮都在当前最好的 agent 上改，立刻有效的改进一个个被加上。")
    print("    但它很快就卡住了：\"多次尝试\"\"另一个模型挑选\"\"并行运行测试\"单独加上都会让得分明显下降，")
    print("    贪心几乎从不接受它们（只有评估噪声偶尔让它们侥幸通过），也就到不了它们的组合 —— 最大的那块提升。")
    print("  · DGM 前期慢一些：它的修改分散在整个存档里，很多次修改花在了\"看起来更差\"的 agent 上。")
    print("    但这些更差、能运行的 agent 成了垫脚石：之后从它们出发再加上搭档，就跨过了\"先变差、再变好\"的低谷。")
    print("    DGM 一直在进步，最后接近最优的设计。这是开放式探索（open-endedness）的核心：不只追求眼前的分数。")
    print("  · 修改者不变时（w/o self-improve），修改的质量一直停在初始 agent 的水平：诊断不准、常常改坏，")
    print("    始终落后于完整的 DGM。完整的 DGM 里，agent 变强 → 修改自己的能力也变强 —— 自我改进的复利。")
    print("  · DGM 论文的消融结果方向相同：两个消融版本都明显不如完整的 DGM（SWE-bench 上 DGM 从 20.0% 提高到 50.0%）。")
    print("    这里的数字完全取决于设定的得分函数和修改质量；编写时试过更浅的低谷（单独的代价只有 −0.04），")
    print("    那时评估噪声足以让贪心跨过低谷，贪心反而最好。**开放式探索的价值取决于问题有多\"欺骗性\"。**\n")


if __name__ == "__main__":
    main()