"""
=============================================================================
第 10 讲 · 01：CodeMonkeys —— 串行 + 并行的 test-time compute，用模型写的测试来挑选
=============================================================================

对应讲义：lectures/lecture10_swe_agents.md 第 2 节（Ehrlich et al. 2025）

-----------------------------------------------------------------------------
【toy world 里的"修 issue"】
-----------------------------------------------------------------------------
一个 issue = 第 1 讲的一道算术题（"正确的修复"= 正确答案）。10B 模型（temperature = 0.7）扮演 coding agent。
每条**轨迹**（trajectory）：
    1. 测试状态机：模型先写一个"复现测试"—— 它自己算一遍，得到它认为正确的答案 t（测试断言 "输出 == t"）。
       模型会算错，所以**测试也会写错**；而且测试和修复出自同一个模型，错误是相关的。
    2. 编辑状态机：写一个修复（一个解答），运行自己的测试；不通过就修改（第 2 讲的 revise），最多 L 轮。
       （论文里两个状态机会互相调整："测试本身写错了"时也会去改测试。这里简化为测试写好就不再改。）
    3. 这条轨迹的最终修复成为一个**候选**。
并行地跑 T 条轨迹，得到 T 个候选、T 个测试。串行（L）和并行（T）是两种扩展 test-time compute 的方式。

挑选（"真正的测试"是隐藏的，只用来评估）：
    · 随机挑一个候选；
    · 测试投票：每个候选在**所有** T 个模型写的测试上跑一遍，选通过最多的；
    · 投票 + 选择状态机：投票的前 3 名交给一个更强的模型（100B）做最终选择（它自己仔细地算一遍，
      选和它的结果一致的候选；都不一致就选投票第一的）；
    · oracle：只要有一个候选是对的就算成功（= coverage，挑选的上限）。

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

import os
import random
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(_HERE, "..", "lecture02"))

from tts_tools import ToyLLM, make_problem_set, mean, pad, revise  # noqa: E402

TEMPERATURE = 0.7


def trajectory(model, problem, rng, max_iters, temperature=TEMPERATURE):
    """
    一条轨迹：写测试 → 写修复 → 跑测试 → 不通过就修改，最多 max_iters 轮。

    【返回】(最终的修复, 这条轨迹写的测试的期望值)
    【边界】测试写错时，"通过测试"反而意味着修复是错的 —— 模型会朝着错误的测试去修改。
    """
    test_value = model.generate(problem, rng, temperature).answer        # 模型认为正确的答案
    edit = model.generate(problem, rng, temperature)
    for _ in range(max_iters - 1):
        if edit.answer == test_value:
            break
        edit = revise(model, problem, edit, rng, temperature)
    return edit, test_value


def run_system(model, problem, rng, n_traj, max_iters, temperature=TEMPERATURE):
    """并行地跑 n_traj 条轨迹，返回 (候选列表, 测试列表)。"""
    outs = [trajectory(model, problem, rng, max_iters, temperature) for _ in range(n_traj)]
    return [e for e, _ in outs], [t for _, t in outs]


def select(problem, candidates, tests, rng, method, selector=None):
    """
    从候选里挑一个，返回是否选中了正确的修复。

    【参数】method："random" / "vote" / "vote+selector" / "oracle"；selector：选择状态机用的模型
    【思路】测试投票：候选在所有模型写的测试上的通过数 —— 本质上是"和多少条轨迹的测试一致"。
        如果模型最常犯的错误让很多测试写成了同一个错误的值，投票就会偏向那个错误的修复（第 1 讲 majority vote 的问题）。
    """
    if method == "oracle":
        return any(c.correct for c in candidates)
    if method == "random":
        return candidates[rng.randrange(len(candidates))].correct
    votes = [sum(1 for t in tests if c.answer == t) for c in candidates]
    order = sorted(range(len(candidates)), key=lambda i: -votes[i])
    if method == "vote":
        return candidates[order[0]].correct
    top = [candidates[i] for i in order[:3]]
    careful = selector.generate(problem, rng, 0.3).answer                # 选择者仔细地算一遍
    for c in top:
        if c.answer == careful:
            return c.correct
    return top[0].correct


def main():
    model, selector = ToyLLM(1e10), ToyLLM(1e11)
    problems = make_problem_set(random.Random(220), n_problems=150, min_steps=4, max_steps=8)
    rng = random.Random(221)

    print("=" * 78)
    print(f"演示 1：coverage（至少一个候选是对的）—— 并行的轨迹数 T × 每条轨迹的最大轮数 L（{model.name} 模型，150 个 issue）")
    print("=" * 78)
    print(pad("T \\ L", 8) + "".join(pad(f"L = {l}", 10, "right") for l in (1, 2, 4, 8)))
    for t in (1, 2, 5, 10):
        row = pad(f"T = {t}", 8)
        for l in (1, 2, 4, 8):
            row += pad(f"{mean(any(c.correct for c in run_system(model, p, rng, t, l)[0]) for p in problems):.3f}",
                       10, "right")
        print(row)
    print("\n观察：")
    print("  · 并行（更多轨迹）和串行（每条轨迹更多轮的修改）都能提高 coverage，两者可以组合。")
    print("  · 串行的收益来得快、但很快饱和：修改是朝着模型**自己写的测试**去改的，测试写错了，越改越偏；")
    print("    并行的收益更持久：不同的轨迹写出不同的测试、不同的修复，覆盖了更多的可能。")
    print("  · CodeMonkeys 在 SWE-bench Verified 上用 10 条轨迹 × 最多 8 轮，coverage 达到 69.8%。\n")

    print("=" * 78)
    print("演示 2：挑选 —— T = 10 条轨迹，从 10 个候选里挑一个；每条轨迹不迭代（L = 1）或迭代 4 轮（L = 4）")
    print("=" * 78)
    methods = (("随机挑一个", "random"), ("模型写的测试投票", "vote"),
               (f"投票 + 选择状态机（{selector.name}）", "vote+selector"), ("oracle（coverage）", "oracle"))
    print(pad("挑选方法", 34) + pad("L = 1", 10, "right") + pad("L = 4", 10, "right"))
    results = {}
    for l in (1, 4):
        systems = {p.pid: run_system(model, p, rng, 10, l) for p in problems}
        results[l] = {m: mean(select(p, *systems[p.pid], rng, m, selector) for p in problems) for _, m in methods}
    for name, m in methods:
        print(pad(name, 34) + pad(f"{results[1][m]:.3f}", 10, "right") + pad(f"{results[4][m]:.3f}", 10, "right"))
    for l in (1, 4):
        r = results[l]
        gap = (r["vote+selector"] - r["random"]) / (r["oracle"] - r["random"])
        print(f"  L = {l}：投票 + 选择状态机 收回了\"随机\"到\"oracle\"之间差距的 {gap:.0%}")
    print("\n观察：")
    print("  · 不迭代时（L = 1），模型写的测试投票比随机好得多，再加上更强的选择者又好一截，收回了大部分差距。")
    print("  · 迭代 4 轮之后，随机挑一个已经不差了，**投票反而比随机还差**：每条轨迹都朝着**自己的**测试修改过，")
    print("    候选和测试高度一致。模型最常犯的那个错误，同时占据了很多候选和很多测试，投票就会偏向它；")
    print("    选择者也只在投票的前 3 名里挑，帮不上忙。测试里的信息在编辑时已经被用掉了，而且是以\"相关\"的方式。")
    print("  · 两种情况下离 coverage 都还有很大的差距：测试和修复出自同一个模型，它最常犯的错误会同时出现在")
    print("    测试和修复里，互相\"确认\"（第 1 讲 self-check 的问题）。更强的选择者能收回一部分差距 ——")
    print("    更强的 verifier 缩小了 generation–verification gap（第 3 讲）。")
    print("  · CodeMonkeys 报告，它的挑选方法收回了随机和 oracle 之间大约一半的差距（coverage 69.8%，最终 57.4%）；")
    print("    它的选择状态机还能**写新的测试**来区分候选，比这里\"自己算一遍\"的选择者强得多。\n")

    print("=" * 78)
    print("演示 3：把几个实力相当的系统的候选合在一起挑选（论文里的 Barrel of Monkeys）")
    print("=" * 78)
    others = (("A：30B，温度 0.7", ToyLLM(3e10), 0.7), ("B：30B，温度 1.0", ToyLLM(3e10), 1.0),
              ("C：20B，温度 0.7", ToyLLM(2e10), 0.7))
    per_system, pooled = {name: [] for name, _, _ in others}, []
    for p in problems:
        all_c, all_t = [], []
        for name, m, temp in others:
            c, t = run_system(m, p, rng, 10, 4, temp)
            per_system[name].append(select(p, c, t, rng, "vote+selector", selector))
            all_c.extend(c)
            all_t.extend(t)
        pooled.append(select(p, all_c, all_t, rng, "vote+selector", selector))
    for name, _, _ in others:
        print(f"  系统 {pad(name, 20)} 单独挑选：{mean(per_system[name]):.3f}")
    print(f"  三个系统的 30 个候选合在一起挑选：{mean(pooled):.3f}")
    print("\n观察：")
    print("  · 在这个 toy world 里，合在一起挑选**没有**超过最好的单个系统。三个系统是同一个\"模型家族\"，")
    print("    共享同样的系统性错误；合并之后，这些共同的错误在测试投票里更占优势，正确的候选反而更难进入前 3 名。")
    print("  · CodeMonkeys 把 SWE-bench Verified 排行榜上几个顶级系统（来自不同的团队、不同的方法）的候选合在一起，")
    print("    得到 66.2%，超过了其中最好的单个系统。差别在于：那些系统**错得不一样**，而且选择 agent 很强。")
    print("    合并候选的价值取决于多样性和挑选的能力 —— 这又是\"错得不一样\"的原则（第 1、3 讲）。\n")


if __name__ == "__main__":
    main()