"""
=============================================================================
第 2 讲 · 工具模块：test-time scaling 的几个积木 —— 修改（revise）、verifier、融合（fuse）
=============================================================================

对应讲义：lectures/lecture02_test_time_compute.md

本模块建立在第 1 讲的 toy world（code/lecture01/toy_llm.py）之上，增加三种 test-time 的"操作"：

1. revise（顺序修改，sequential revision）
   模型重读自己上一次的解答，**一步一步地复查**：每一步用自己的能力、以较低的温度重新算一遍，
   如果重算的结果和写下的不一样，就认为这一步错了，用重算的值替换它，并从这里往后重新写。
   这个机制没有任何新参数，却自然地产生了三个真实的现象：
       · 随机手滑（careless）容易被发现：重算一遍通常得到不同的值；
       · 固定的错误套路（bug）很难被发现：如果模型"最相信"这个套路，仔细重算也会再犯一次同样的错；
       · 修改可能把对的改错：原本正确的一步，重算时算错了，反而被"纠正"成错的。
   每次修改最多纠正一处（找到第一处不一致就从那里重写）。所以"只错了一步"的简单题，
   一两次修改就能修好；错误多、或者错在模型自己的盲点上的难题，修改帮助不大。

2. verifier_score（打分模型）
   和第 1 讲 02 号文件演示 3 一样：正确答案平均高 signal 分，再加上标准正态噪声；
   噪声按"(题号, 答案)"固定 —— 同一个错误答案每次得到同样的分数（verifier 有自己固定的误解）。

3. fuse（融合，fusion）
   Archon 里的 Fuser：读入多个候选解答，综合成一个新的解答。这里的做法是"逐步投票"：
   从第一步开始，看和目前融合结果"前缀一致"的候选在这一步写了什么值，多数一致就采纳，
   否则融合模型自己算。它能把"各自错在不同地方"的候选拼成一个全对的解答 ——
   这是只在最终答案上投票（majority vote）做不到的。

-----------------------------------------------------------------------------
【用法】
-----------------------------------------------------------------------------
    import random
    from tts_tools import ToyLLM, make_problem, revise, verifier_score

    rng = random.Random(0)
    problem = make_problem(rng, 5)
    model = ToyLLM(1e10)
    s0 = model.generate(problem, rng)
    s1 = revise(model, problem, s0, rng)        # 修改一次
    print(s0.answer, s1.answer, problem.answer())
"""

import math
import os
import random
import sys

# 【Python 小知识】sys.path 是 Python 查找模块的目录列表。把第 1 讲的目录插到最前面，
# 就可以直接 `from toy_llm import ...`。os.path.abspath(__file__) 是本文件的绝对路径，
# 这样无论你从哪个目录运行脚本，都能找到第 1 讲。
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(_HERE, "..", "lecture01"))

from toy_llm import (CARELESS_OFFSETS, PROCEDURES, Problem, Sample, ToyLLM, apply_bug,  # noqa: E402
                     apply_op, bar, fmt_size, majority_vote, make_problem, make_problem_set, mean,
                     pad, pass_at_k, sample_index)

# 【Python 小知识】__all__ 列出"from tts_tools import *"时导出的名字；这里也起到"本模块提供什么"的说明作用。
__all__ = ["CARELESS_OFFSETS", "PROCEDURES", "Problem", "Sample", "ToyLLM", "apply_bug", "apply_op", "bar",
           "fmt_size", "majority_vote", "make_problem", "make_problem_set", "mean", "pad", "pass_at_k",
           "sample_index", "sample_step", "extend_trace", "revise", "run_chain", "verifier_score",
           "select_by_verifier", "select_weighted", "fuse", "estimate_pass1", "difficulty_bins", "CHECK_TEMPERATURE"]


# =============================================================================
# 第一部分：一步一步地生成（第 1 讲 generate 的拆分版）
# =============================================================================
def sample_step(model, step, x, step_index, rng, temperature=1.0):
    """
    让模型做一步运算：按概率选一种做法，算出这一步的值。

    【思路】和 toy_llm.ToyLLM.generate 里每一步的逻辑完全相同，单独拿出来，
        是为了在"复查""续写""融合"里复用。
    【参数】model：ToyLLM；step：Step；x：这一步的输入值；step_index：第几步（从 0 开始）；
        rng：random.Random；temperature：采样温度
    【返回】(新值, 做法名)
    """
    probs = model.step_probs(step.op, x, step.operand, step_index, temperature)
    proc = PROCEDURES[sample_index(probs, rng)]
    right = apply_op(step.op, x, step.operand)
    if proc == "correct":
        return right, proc
    if proc == "bug":
        return apply_bug(step.op, x, step.operand), proc
    return right + CARELESS_OFFSETS[rng.randrange(len(CARELESS_OFFSETS))], proc


def extend_trace(model, problem, values, procs, rng, temperature=1.0):
    """
    从已经写好的前几步出发，把解答续写完，返回一个完整的 Sample。

    【参数】
        model, problem, rng, temperature：同上
        values：已经写好的各步的值（list，可以为空 = 从头写）
        procs ：这些步用的做法（list，和 values 等长）
    【返回】Sample（local_ok 按"每一步就它自己的输入而言是否算对"重新计算）
    【边界】
        · 不会修改传进来的 values、procs（先复制一份）。
        · len(values) == problem.n_steps 时不再生成，直接打包返回。
    【自我改进联系】第 3 讲的 Math-Shepherd 正是用"从某一步续写很多次"来估计这一步之后还能不能做对。
    """
    values, procs = list(values), list(procs)
    x = values[-1] if values else problem.start
    for i in range(len(values), problem.n_steps):
        x, proc = sample_step(model, problem.steps[i], x, i, rng, temperature)
        values.append(x)
        procs.append(proc)
    return _pack(problem, values, procs)


def _pack(problem, values, procs):
    """把各步的值和做法打包成 Sample，顺便算出 local_ok 和 correct。（内部函数，名字以 _ 开头）"""
    local_ok = []
    x = problem.start
    for step, v in zip(problem.steps, values):
        local_ok.append(v == apply_op(step.op, x, step.operand))
        x = v
    answer = values[-1]
    return Sample(answer=answer, correct=(answer == problem.answer()), values=tuple(values),
                  procs=tuple(procs), local_ok=tuple(local_ok))


# =============================================================================
# 第二部分：顺序修改（sequential revision）
# =============================================================================
CHECK_TEMPERATURE = 0.3   # 复查时的温度：比生成时低，表示"仔细地"重算


def revise(model, problem, sample, rng, temperature=1.0, check_temperature=CHECK_TEMPERATURE):
    """
    让模型复查并修改自己的一次解答（最多纠正一处）。

    【思路】见文件开头。逐步复查：第 i 步用自己的能力、以较低的温度（仔细地）重算一次，
        · 重算结果和写下的值相同 → 认为这一步没问题，继续看下一步；
        · 不同 → 认为这一步错了，用重算的值替换它，再从第 i+1 步开始（以正常温度）重新往后写，结束这次修改。
        如果每一步都"复查通过"，返回原来的解答（模型觉得没什么可改的）。
        为什么复查要用低温？如果复查和生成一样"随手"，正确的步骤也会经常被误报成错误，
        修改就退化成"随机重写后半段"。低温复查更接近模型"最相信的答案"：
        它能发现手滑（careless），却会坚持自己的误解（bug 的 logit 比 correct 高的那些步骤）。
        所以反复修改的结果会**收敛到模型自己最相信的解答**附近 —— 修改是在"利用"模型已有的知识，
        而不是"探索"它不知道的东西。
    【参数】model；problem；sample：上一次的解答（Sample）；rng；
        temperature：续写时的温度；check_temperature：复查时的温度（越低越"固执"）
    【返回】新的 Sample（可能和原来的相同）
    【边界】
        · 复查用的是模型自己的能力：它对自己的 systematic error 基本是"盲"的，
          这和第 1 讲 03 号文件里 self-check 的局限是同一个原因。
        · 原本正确的一步，复查时仍可能（以较小的概率）算错，于是被"改错" —— 修改不一定单调变好。
        · 每次修改的成本记为"一次生成"（复查 + 续写的 token 数和重新写一遍差不多）。
    【自我改进联系】Snell et al. 2024 专门微调了一个"会修改"的模型（revision model）；
        这里的机制是一个卡通，用来说明"修改擅长修小错、不擅长修根本性的错"。
    """
    x = problem.start
    for i, step in enumerate(problem.steps):
        recomputed, proc = sample_step(model, step, x, i, rng, check_temperature)
        if recomputed != sample.values[i]:
            values = list(sample.values[:i]) + [recomputed]
            procs = list(sample.procs[:i]) + [proc]
            return extend_trace(model, problem, values, procs, rng, temperature)
        x = sample.values[i]
    return sample


def run_chain(model, problem, rng, length, temperature=1.0):
    """
    一条修改链：先生成一次，再连续修改 length − 1 次。返回链上全部 length 个解答。

    【参数】length：链长（≥ 1）。length = 1 就是普通的一次采样。
    【返回】Sample 的列表，第 0 个是初始解答，第 j 个是修改 j 次之后的解答
    【边界】length < 1 抛出 ValueError。
    """
    if length < 1:
        raise ValueError("length 必须 ≥ 1")
    chain = [model.generate(problem, rng, temperature)]
    for _ in range(length - 1):
        chain.append(revise(model, problem, chain[-1], rng, temperature))
    return chain


# =============================================================================
# 第三部分：verifier 与挑选
# =============================================================================
def verifier_score(pid, answer, correct, signal=2.0, seed=0):
    """
    一个有噪声的 outcome verifier 给"第 pid 题的答案 answer"打的分数。

    【思路】分数 = signal × [答案正确] + 标准正态噪声。噪声由 (seed, 题号, 答案) 唯一决定：
        同一个答案无论出现多少次、以什么顺序被查询，分数都一样 —— 模拟 verifier 自己"固定的误解"。
        signal = 2.0 时，随便拿一个正确答案和一个错误答案，verifier 排对的概率约 92%。
    【参数】pid：题号；answer：答案；correct：答案是否正确（上帝视角，只用来模拟 verifier 的准确率）；
        signal：区分度；seed：换一个 seed 相当于换一个 verifier
    【返回】float
    【Python 小知识】random.Random("字符串")：用字符串做种子。字符串会经过 SHA-512 变成整数，
        结果不受 PYTHONHASHSEED 影响，在不同 Python 版本、不同运行之间都相同。
    """
    noise = random.Random(f"{seed}|{pid}|{answer}").gauss(0.0, 1.0)
    return (signal if correct else 0.0) + noise


def select_by_verifier(problem, samples, signal=2.0, seed=0):
    """
    best-of-N：在候选里选 verifier 分数最高的那个。

    【参数】problem；samples：Sample 列表（非空）；signal、seed：传给 verifier_score
    【返回】被选中的 Sample
    【边界】samples 为空时 max 会抛出 ValueError。分数相同（同一个答案）时取最先出现的。
    """
    return max(samples, key=lambda s: verifier_score(problem.pid, s.answer, s.correct, signal, seed))


def select_weighted(problem, samples, signal=2.0, seed=0, sharpness=2.0):
    """
    weighted best-of-N（Snell et al. 2024、Lightman et al. 2023 都用这种挑法）：
    把同一个答案的所有样本的 verifier 概率**加起来**，选总分最高的答案。

    【思路】先把 verifier 的分数变成"这个答案正确的概率"：prob = σ(sharpness × (分数 − signal/2))，
        分数恰好在"正确"和"错误"两类的中点时 prob = 0.5。然后按答案求和：
            总分(答案 a) = Σ_{答案为 a 的样本} prob(a) = 出现次数(a) × prob(a)
        它同时用到了两种信号：**出现次数**（majority vote）和 **verifier 的判断**。
        为什么比"只看 verifier 最高分"稳健？候选很多时，总有某个错误答案"运气好"拿到很高的分数；
        但如果它只出现一两次，乘上出现次数之后就比不过反复出现、分数也不错的正确答案。
    【参数】problem；samples：Sample 列表（非空）；signal、seed：传给 verifier_score；sharpness：σ 的陡峭程度
    【返回】一个答案为"总分最高的答案"的 Sample（取最先出现的那个样本）
    【边界】如果模型最常犯的错误答案同时也被 verifier 看好，weighted 也会选错 —— 它继承了 majority vote
        的弱点（systematic error），只是比单纯投票多了 verifier 的纠正。
    """
    totals = {}
    first = {}
    for s in samples:
        score = verifier_score(problem.pid, s.answer, s.correct, signal, seed)
        prob = 1.0 / (1.0 + math.exp(-sharpness * (score - signal / 2)))
        totals[s.answer] = totals.get(s.answer, 0.0) + prob
        first.setdefault(s.answer, s)        # 【Python 小知识】setdefault：键不存在时才写入
    best = max(totals, key=lambda a: totals[a])
    return first[best]


# =============================================================================
# 第四部分：融合（fusion）
# =============================================================================
def fuse(model, problem, candidates, rng, temperature=1.0):
    """
    Archon 式的 Fuser：读入多个候选解答，逐步投票，综合成一个新的解答。

    【思路】融合模型从第一步开始往后写。写第 i 步时：
        1. 只看和"目前融合出的前 i 步"完全一致的候选（前缀一致，它们第 i 步的输入和我们相同）；
        2. 如果这些候选里，某个值得票严格过半（> 一半），就采纳它（做法记为 "fused"）；
        3. 否则（没有一致的候选、或者意见分散），融合模型自己算这一步。
        和"只在最终答案上投票"相比，逐步投票能把"各自错在不同步骤"的候选拼成一个全对的解答。
    【参数】model：融合用的模型（可以和生成候选的模型不同）；problem；candidates：Sample 列表；rng；temperature
    【返回】Sample
    【边界】
        · 候选为空时，等于融合模型自己从头写一遍。
        · 如果多数候选在同一步犯了同一个 systematic error，融合会采纳这个错误 —— 投票救不了"大家都错得一样"。
        · 真实的 Fuser 是一个读了所有候选的 LLM，能力远比"逐步投票"复杂；这里只保留"综合信息"这一核心。
    """
    values, procs = [], []
    x = problem.start
    for i, step in enumerate(problem.steps):
        consistent = [c for c in candidates if list(c.values[:i]) == values]
        counts = {}
        for c in consistent:
            counts[c.values[i]] = counts.get(c.values[i], 0) + 1
        best_value, best_count = None, 0
        for v, cnt in counts.items():
            if cnt > best_count:
                best_value, best_count = v, cnt
        if consistent and best_count * 2 > len(consistent):
            x, proc = best_value, "fused"
        else:
            x, proc = sample_step(model, step, x, i, rng, temperature)
        values.append(x)
        procs.append(proc)
    return _pack(problem, values, procs)


# =============================================================================
# 第五部分：难度
# =============================================================================
def estimate_pass1(model, problem, rng, n, temperature=1.0):
    """
    用 n 次采样估计一道题的单次成功率 pass@1。

    【边界】n 太小时估计很粗：n = 32 时，最小的非零估计是 1/32 ≈ 3%。
    """
    return mean(model.generate(problem, rng, temperature).correct for _ in range(n))


def difficulty_bins(pass1_values, n_bins=5):
    """
    按 pass@1 把题目分成 n_bins 个难度档（Snell et al. 2024 的做法：按分位数等分）。

    【思路】把题目按 pass@1 从高到低排序，再平均切成 n_bins 段：第 0 档最简单，最后一档最难。
    【参数】pass1_values：每道题的 pass@1 估计（list）；n_bins：档数
    【返回】list，第 i 个元素是第 i 道题的档号（0 ~ n_bins−1）
    【边界】pass@1 相同的题按题目顺序排（排序是稳定的），所以并列的题可能被分到相邻的两档。
    """
    order = sorted(range(len(pass1_values)), key=lambda i: -pass1_values[i])
    bins = [0] * len(pass1_values)
    for rank, i in enumerate(order):
        bins[i] = rank * n_bins // len(pass1_values)
    return bins


if __name__ == "__main__":
    # 小演示：一道题，看看修改链和融合的效果
    rng = random.Random(0)
    model = ToyLLM(1e10)
    problem = make_problem(rng, 6)
    print(problem.text(), "标准答案：", problem.answer())
    chain = run_chain(model, problem, rng, length=5)
    for j, s in enumerate(chain):
        print(f"  修改 {j} 次：答案 {s.answer:>8d}  {'✓' if s.correct else '✗'}  做法 {', '.join(s.procs)}")
    cands = [model.generate(problem, rng) for _ in range(5)]
    fused = fuse(model, problem, cands, rng)
    print("  5 个候选的答案：", [c.answer for c in cands], " 融合后：", fused.answer, "✓" if fused.correct else "✗")