"""
=============================================================================
第 4 讲 · 02：RLEF —— 用强化学习教模型利用执行反馈
=============================================================================

对应讲义：lectures/lecture04_feedback_tools_code.md 第 3 节（Gehring et al. 2024，RLEF）

-----------------------------------------------------------------------------
【问题】
-----------------------------------------------------------------------------
写代码时，运行测试会得到反馈：哪个测试没过、报了什么错。直觉上，模型应该能利用这些反馈修改代码。
但 RLEF 论文发现：**现有的 LLM 很不擅长利用执行反馈**——多轮"根据反馈修改"的效果，
往往还不如同样次数的"独立地重新采样"。能不能用强化学习教会它们？

-----------------------------------------------------------------------------
【这里的设定】
-----------------------------------------------------------------------------
每道题最多 3 轮。每一轮：
    1. 模型给出一个解答（第 1 轮从头生成；之后按下面的"策略"决定怎么写）；
    2. "运行测试"：如果答案对，结束（奖励 1）；否则得到反馈："第 i 步的结果不对（你写的是 v）"——
       i 是第一个出错的步骤，就像测试报告里指出了第一个失败的断言。
模型在第 2、3 轮有两种做法：
    · 利用反馈（repair）：保留第 i 步之前的内容，在第 i 步**排除已知错误的值 v** 再重算，然后往后续写。
      "排除"很关键：如果错误来自模型"最相信"的错误套路，从头重采样很可能再犯一次；
      而反馈明确告诉它"这个值是错的"，它就不会再写出同样的值。
    · 忽略反馈（resample）：从头重新生成（等于独立采样）。
策略只有一个参数 θ：利用反馈的概率 = σ(θ)。初始时 σ(θ) = 0.1 —— 模型基本上不理会反馈。
用 REINFORCE 训练 θ：奖励 = 3 轮之内是否做对。

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

import math
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 (CARELESS_OFFSETS, PROCEDURES, ToyLLM, apply_bug, apply_op, extend_trace,  # noqa: E402
                       make_problem_set, mean, pad, sample_index)

TEMPERATURE = 0.7
MAX_TURNS = 3
THETA0 = math.log(0.1 / 0.9)      # σ(θ0) = 0.1


def sigmoid(z):
    return 1 / (1 + math.exp(-z)) if z >= 0 else math.exp(z) / (1 + math.exp(z))


def first_wrong_step(sample):
    """执行反馈：第一个算错的步骤（局部检查）。答案对但过程错时可能返回某个下标——这里只在答案错时调用。"""
    for i, ok in enumerate(sample.local_ok):
        if not ok:
            return i
    return len(sample.local_ok) - 1   # 极少数情况：每一步局部都对、答案却错（不会发生在本 toy world 里，保险起见）


def repair(model, problem, sample, i, rng, avoid_value):
    """
    利用反馈修改：保留前 i 步，第 i 步排除已知错误的值 avoid_value 再算一次，然后续写。

    【思路】三种做法里，凡是"一定会再写出 avoid_value"的做法（例如 bug 的结果恰好就是它）概率置 0，
        剩下的重新归一化后抽一个；careless 的偏移也避开 avoid_value。
    【参数】i：要修改的步骤；avoid_value：反馈指出的错误值
    【返回】新的 Sample
    【边界】这一步之后的步骤从头续写，仍然可能出新的错 —— 下一轮的反馈会再指出来。
    """
    step = problem.steps[i]
    x = problem.start if i == 0 else sample.values[i - 1]
    probs = model.step_probs(step.op, x, step.operand, i, TEMPERATURE)
    right = apply_op(step.op, x, step.operand)
    bug = apply_bug(step.op, x, step.operand)
    allowed = [p if not ((proc == "correct" and right == avoid_value) or (proc == "bug" and bug == avoid_value))
               else 0.0 for proc, p in zip(PROCEDURES, probs)]
    total = sum(allowed)
    proc = PROCEDURES[sample_index([a / total for a in allowed], rng)]
    if proc == "correct":
        v = right
    elif proc == "bug":
        v = bug
    else:
        offsets = [d for d in CARELESS_OFFSETS if right + d != avoid_value]
        v = right + offsets[rng.randrange(len(offsets))]
    return extend_trace(model, problem, list(sample.values[:i]) + [v], list(sample.procs[:i]) + [proc],
                        rng, TEMPERATURE)


def episode(model, problem, theta, rng, feedback="true"):
    """
    跑一个最多 MAX_TURNS 轮的回合。

    【参数】feedback："true"（指出第一个错误）或 "random"（随机指一步 —— 消融实验：没有信息的反馈）
    【返回】(奖励 0/1, 各次决策的 ∂log π/∂θ 之和)
    【思路】∂ log σ(θ) / ∂θ = 1 − σ(θ)（选了"利用反馈"时）；∂ log(1 − σ(θ)) / ∂θ = −σ(θ)（选了"忽略"时）。
    """
    p_use = sigmoid(theta)
    grad = 0.0
    sample = model.generate(problem, rng, TEMPERATURE)
    for _turn in range(1, MAX_TURNS):
        if sample.correct:
            return 1, grad
        i = first_wrong_step(sample) if feedback == "true" else rng.randrange(problem.n_steps)
        if rng.random() < p_use:
            grad += 1 - p_use
            sample = repair(model, problem, sample, i, rng, sample.values[i])
        else:
            grad -= p_use
            sample = model.generate(problem, rng, TEMPERATURE)
    return (1 if sample.correct else 0), grad


def evaluate(model, problems, theta, rng, feedback="true", repeats=3):
    """在一组题上评估：3 轮之内做对的比例。"""
    return mean(episode(model, p, theta, rng, feedback)[0] for p in problems for _ in range(repeats))


def train(model, problems, rng, feedback="true", epochs=40, lr=20.0):
    """
    REINFORCE 训练 θ：每个 epoch 把训练集跑一遍，用"奖励 − 基线"乘以 ∂log π/∂θ 的平均值更新 θ。

    【思路】基线（baseline）= 上一个 epoch 的平均奖励。减去基线不改变梯度的期望，但能大大减小方差：
        只有"比平时好"的回合才会推动 θ 往它选择的方向走（第 6 讲会详细讲）。
        学习率为什么这么大（20）？一开始模型只有 10% 的概率选择"利用反馈"，这个选择很少出现，
        梯度的期望每个 epoch 只有 0.01 左右。学习率太小，几十个 epoch 也几乎学不动（编写时试过 lr = 1）。
        这也是 RL 的一个普遍问题：**初始策略很少尝试的行为，很难被学会**（探索不足）。
    【返回】每 5 个 epoch 记录一次的 (epoch, σ(θ), 训练集上的平均奖励)
    """
    theta, baseline, log = THETA0, None, []
    for epoch in range(1, epochs + 1):
        results = [episode(model, p, theta, rng, feedback) for p in problems]
        rewards = [r for r, _ in results]
        if baseline is None:
            baseline = mean(rewards)
        theta += lr * mean((r - baseline) * g for r, g in results)
        baseline = mean(rewards)
        if epoch % 5 == 0 or epoch == 1:
            log.append((epoch, sigmoid(theta), baseline))
    return theta, log


def main():
    model = ToyLLM(1e10)
    train_set = make_problem_set(random.Random(100), n_problems=200, min_steps=3, max_steps=8)
    test_set = make_problem_set(random.Random(101), n_problems=200, min_steps=3, max_steps=8)

    print("=" * 78)
    print(f"演示 1：RL 训练\"是否利用执行反馈\"（{model.name} 模型，最多 {MAX_TURNS} 轮，200 道训练题）")
    print("=" * 78)
    theta, log = train(model, train_set, random.Random(102))
    print(pad("epoch", 8) + pad("利用反馈的概率 σ(θ)", 22, "right") + pad("训练集上 3 轮内做对", 22, "right"))
    for epoch, p_use, reward in log:
        print(pad(str(epoch), 8) + pad(f"{p_use:.3f}", 22, "right") + pad(f"{reward:.3f}", 22, "right"))

    rng = random.Random(103)
    before = evaluate(model, test_set, THETA0, rng)
    after = evaluate(model, test_set, theta, rng)
    always = evaluate(model, test_set, 20.0, rng)        # σ(20) ≈ 1：总是利用反馈
    # 独立采样 + 测试：k 次里有一次通过就算做对（测试就是 oracle）
    indep = {k: mean(any(model.generate(p, rng, TEMPERATURE).correct for _ in range(k))
                     for p in test_set for _ in range(3)) for k in (1, 3, 5, 10, 20)}
    print("\n测试集（另外 200 道题）上，3 轮之内做对的比例：")
    print(f"  训练前（σ(θ) = 0.1，几乎不理会反馈）：{before:.3f}")
    print(f"  训练后（σ(θ) = {sigmoid(theta):.3f}）：{after:.3f}")
    print(f"  总是利用反馈（σ(θ) = 1）：{always:.3f}")
    print("  对照：独立采样 k 次（用测试挑出对的）：" + "，".join(f"k = {k}：{v:.3f}" for k, v in indep.items()))
    print("\n观察：")
    print("  · 训练前，3 轮多次尝试的效果和\"独立采样 3 次\"差不多：模型基本上不理会反馈，每一轮都是从头重来。")
    print("    这正是 RLEF 论文观察到的现象：多轮修改并不比独立采样好。")
    print("  · RL 只用\"最后有没有做对\"作为奖励，却学会了利用反馈：σ(θ) 在第 5~10 个 epoch 之间迅速从 0.25 升到 0.9，")
    print("    最后接近 1。利用反馈的回合更常做对，REINFORCE 就把\"利用反馈\"的概率推了上去。")
    print("    一开始上升得很慢：这个行为本来很少被尝试，信号很弱（探索不足的问题，第 6 讲还会遇到）。")
    print("  · 训练后，3 轮的效果相当于独立采样 5 次左右（0.843 vs k = 5 的 0.833）：反馈告诉模型\"这个值是错的\"，")
    print("    让它跳出了自己\"最相信\"的错误套路，这是从头重采样做不到的。")
    print("    提升在 toy world 里不算大：每一轮只修一处错误，后面的步骤重写时又可能出新的错。")
    print("    RLEF 论文在竞赛编程上报告的是：训练后的模型用少一个数量级的采样就能达到同样的成绩。\n")

    print("=" * 78)
    print("演示 2：消融 —— 如果反馈没有信息（随机指一步），RL 会学到什么？")
    print("=" * 78)
    theta_r, log_r = train(model, train_set, random.Random(104), feedback="random")
    print(pad("epoch", 8) + pad("利用反馈的概率 σ(θ)", 22, "right") + pad("训练集上 3 轮内做对", 22, "right"))
    for epoch, p_use, reward in log_r:
        print(pad(str(epoch), 8) + pad(f"{p_use:.3f}", 22, "right") + pad(f"{reward:.3f}", 22, "right"))
    rng = random.Random(105)
    print(f"\n测试集上（随机反馈）：训练前 {evaluate(model, test_set, THETA0, rng, 'random'):.3f}，"
          f"训练后 {evaluate(model, test_set, theta_r, rng, 'random'):.3f}，"
          f"总是\"利用\"随机反馈 {evaluate(model, test_set, 20.0, rng, 'random'):.3f}")
    print("\n观察：")
    print("  · 随机的\"反馈\"常常指向一个本来正确的步骤，\"排除这个值\"反而逼着模型写错。")
    print("    \"总是利用\"随机反馈比不理会它差得多（0.61 左右 vs 0.75 左右）。RL 学会了不理会它：σ(θ) 降到接近 0。")
    print("  · 所以 RL 学到的不是\"总是修改\"这种固定的习惯，而是\"反馈有用时才用\"。RLEF 论文也做了类似的检查：")
    print("    把真实的执行反馈换成随机的，训练后的模型表现明显变差 —— 说明它确实在读反馈的内容。\n")


if __name__ == "__main__":
    main()