"""
=============================================================================
第 13 讲 · 02：AlphaProof 的测试时 RL —— 解不出原题，就先解它的变体
=============================================================================

对应讲义：lectures/lecture13_superhuman_reasoning.md 第 3.4 节（AlphaProof，Nature 2025）

-----------------------------------------------------------------------------
【设定】
-----------------------------------------------------------------------------
一个证明器（prover）对难度为 d 的题，每次尝试成功的概率是
    p(d) = sigmoid(2 × (s − d))
s 是证明器的"能力"。一开始 s = 0；目标题的难度是 D（D 越大，离证明器的能力越远）。
每次尝试的证明都由 Lean **完美地**验证：成功就是真的成功。

证明器可以从成功中学习（RL）：每证出一道难度为 d 的题，s 增加 η × (1 − p(d))。
    · 1 − p(d) 是 REINFORCE 式的"优势"：越难（p 越小）的成功，学到的越多；太容易的题（p ≈ 1）几乎学不到东西。
    · 设定：变体和原题相关，在变体上学到的能力**完全迁移**到原题上（这是乐观的假设）。

AlphaProof 的测试时 RL（test-time RL）：比赛时，用一个语言模型生成目标题的**大量变体**（特例、简化、推广），
难度从很容易到和原题一样难；在这些变体上一边证明、一边学习，直到能证出原题。

比较的做法（相同的尝试次数预算）：
    · 只采样：所有尝试都用在原题上，不学习（= 第 2 讲的重复采样）。
      注意"只在原题上做 RL"和它完全一样：第一次成功之前没有任何可以学习的信号，而第一次成功时题已经解出了。
    · 测试时 RL，只练简单的变体（难度在 0 ~ D/3 之间）；
    · 测试时 RL，变体的难度在 0 ~ D 之间均匀分布；
    · 测试时 RL，自适应课程：每次挑"当前成功率最接近 50%"的变体。
    后三种做法每次尝试都有 10% 的概率直接试原题。

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

import math
import random

ETA = 0.05                  # 学习率：每次成功，能力增加 η × 优势
N_VARIANTS = 20             # 生成的变体数（难度均匀地铺开）
P_TARGET = 0.10             # 每次尝试直接试原题的概率
N_RUNS = 20                 # 每个设定重复多少次（取解出的比例）


def p_success(s, d):
    """
    能力为 s 的证明器，在难度为 d 的题上一次尝试成功的概率：sigmoid(2(s − d))。

    【边界】s − d 很负时，exp 的参数很大；这里先判断符号，避免 math.exp 溢出（OverflowError）。
    """
    z = 2.0 * (s - d)
    if z >= 0:
        return 1.0 / (1.0 + math.exp(-z))
    e = math.exp(z)
    return e / (1.0 + e)


def attempts_until_success(q, rng):
    """
    每次成功概率为 q 的独立尝试，第几次第一次成功（几何分布）。

    【思路】直接一次一次地模拟太慢（q 可能只有 10⁻⁶）。几何分布可以一步抽出来：
        u 在 (0, 1) 上均匀，k = ⌊log(u) / log(1 − q)⌋ + 1。
    【边界】q = 0 时永远不会成功，返回一个巨大的数；用 log1p(−q) 而不是 log(1 − q)，q 很小时更精确。
    【Python 小知识】math.log1p(x) = log(1 + x)，x 接近 0 时比直接算 log(1 + x) 精确得多。
    """
    if q <= 0:
        return 10 ** 18
    u = 1.0 - rng.random()                      # (0, 1]，避免 log(0)
    return int(math.log(u) / math.log1p(-q)) + 1


def run(target_d, method, budget, rng):
    """
    在 budget 次尝试内，能不能证出难度为 target_d 的原题。

    【参数】method："sample" / "easy" / "uniform" / "adaptive"
    【返回】True / False
    【思路】在两次成功之间，能力 s 不变，每次尝试"成功"的概率也不变，所以可以用几何分布直接跳到下一次成功：
        · 一次尝试 = 以 P_TARGET 的概率试原题，否则按课程选一个变体；
        · 这一次尝试成功的总概率 q = Σ（选中它的概率 × 它的成功率）；
        · 跳到下一次成功后，按"选中概率 × 成功率"的比例决定是哪道题成功了；
        · 是原题 → 解出；是变体 → 学习（s 增加），然后继续。
    """
    s, used = 0.0, 0
    if method == "sample":
        return attempts_until_success(p_success(s, target_d), rng) <= budget
    hi = target_d / 3 if method == "easy" else target_d
    variants = [hi * (i + 0.5) / N_VARIANTS for i in range(N_VARIANTS)]
    while True:
        if method == "adaptive":
            chosen = min(variants, key=lambda d: abs(p_success(s, d) - 0.5))
            items = [(P_TARGET, target_d), (1 - P_TARGET, chosen)]
        else:
            items = [(P_TARGET, target_d)] + [((1 - P_TARGET) / N_VARIANTS, d) for d in variants]
        weights = [w * p_success(s, d) for w, d in items]
        q = math.fsum(weights)
        used += attempts_until_success(q, rng)
        if used > budget:
            return False
        u, acc, which = rng.random() * q, 0.0, len(items) - 1
        for i, w in enumerate(weights):
            acc += w
            if u < acc:
                which = i
                break
        if which == 0:
            return True                                    # 原题被证出来了
        d = items[which][1]
        s += ETA * (1.0 - p_success(s, d))                  # 在变体上学习


def main():
    gaps = (2, 4, 6, 8, 10, 12)
    methods = (("只采样（不学习）", "sample"), ("测试时 RL：只练简单的变体", "easy"),
               ("测试时 RL：变体难度均匀分布", "uniform"), ("测试时 RL：自适应课程", "adaptive"))
    for budget in (10 ** 3, 10 ** 4):
        print("=" * 78)
        print(f"预算 {budget:,} 次尝试：解出原题的比例（{N_RUNS} 次重复；D = 原题难度 − 初始能力）")
        print("=" * 78)
        print(" " * 30 + "".join(f"{'D = ' + str(g):>8}" for g in gaps))
        for name, method in methods:
            row = ""
            for g in gaps:
                rng = random.Random(13000 + 100 * g + len(method))
                row += f"{sum(run(g, method, budget, rng) for _ in range(N_RUNS)) / N_RUNS:8.2f}"
            width = sum(2 if ord(ch) > 127 else 1 for ch in name)
            print(name + " " * max(1, 30 - width) + row)
        print()
    print("一次尝试的成功率（初始能力 s = 0）：" + "，".join(f"D = {g}：{p_success(0, g):.1e}" for g in gaps))
    print("\n观察：")
    print("  · 只采样：一次尝试的成功率随难度差 D 指数下降，D = 6 时只有 6 × 10⁻⁶，一万次尝试也只解出 5%。")
    print("    重复采样（第 2 讲）只能把\"偶尔能做对\"变成\"做对\"，够不着离能力太远的题。")
    print("  · 测试时 RL：先证变体、从中学习，原题就从\"几乎不可能\"变成\"可以做到\"。但课程决定了要多少尝试：")
    print("    自适应课程（总是练成功率约 50% 的变体）1000 次尝试就能解出 D = 10；难度均匀分布要一万次；")
    print("    只练简单的变体，一万次也只解出 D = 8 的 80%。")
    print("  · 为什么只练简单的不行：能力远高于题目难度时几乎每次都成功，优势 1 − p 接近 0，几乎学不到东西。")
    print("    学习发生在能力的边界上 —— 这和第 6 讲 DAPO 的\"动态采样\"（去掉全对、全错的题）是同一个道理。")
    print("  · 这一切的前提是 verifier 完美：Lean 保证每个\"成功\"都是真的。如果 verifier 会把错的证明判成对的，")
    print("    测试时 RL 就会在错误上越学越偏。另外，\"变体上的能力完全迁移到原题\"是乐观的设定。\n")


if __name__ == "__main__":
    main()