"""
=============================================================================
第 3 讲 · 工具模块：几种 verifier —— 检查结果（ORM）、检查过程（PRM）、rollout 估值
=============================================================================

对应讲义：lectures/lecture03_verification.md

-----------------------------------------------------------------------------
【toy world 里的 verifier 是什么？】
-----------------------------------------------------------------------------
真实的 verifier 是一个读了解答文本、输出"对不对"的模型。在 toy world 里，我们让 verifier 也是一个 ToyLLM
（称为 checker），它用**自己的能力**去检查。检查的方式不同，就得到不同类型的 verifier：

    · ORM（outcome reward model，检查结果）：checker 不看解答过程，自己从头把题做 m 遍，
      看有多少遍得到了和候选相同的答案。—— 它只能判断"最终答案像不像对的"，
      而且要判断一道 k 步题，checker 自己就要连对 k 步，错误同样会累积。
    · PRM（process reward model，检查过程）：checker 逐步检查候选写下的每一步：
      "从这一步的输入出发，我会以多大的概率写出这个值？"（step_prob）。
      每次只需要判断**一步**，比从头解一遍容易得多；而且能指出**哪一步**可疑。
      整个解答的分数 = 各步分数的乘积（Lightman et al. 2023 的做法）或最小值。
    · rollout 估值（Math-Shepherd）：从候选的第 i 步之后，让一个 completer 模型续写 n 次，
      看有多少次能得到正确答案。这需要知道标准答案（训练时有），用来**自动生成过程标签**。

这些都是卡通：真实的 ORM / PRM 是训练出来的打分模型，不会真的"重算"。但它们刻画了
"检查一步比解完整道题容易""checker 和生成者有同样的盲点"这些真实的机制。

-----------------------------------------------------------------------------
【用法】
-----------------------------------------------------------------------------
    import random
    from verifier_tools import ToyLLM, make_problem, prm_step_scores, aggregate, orm_score

    rng = random.Random(0)
    problem = make_problem(rng, 5)
    gen, checker = ToyLLM(1e10), ToyLLM(1e11)
    s = gen.generate(problem, rng, 0.7)
    print(prm_step_scores(checker, problem, s), aggregate(prm_step_scores(checker, problem, s)))
"""

import math
import os
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, Problem, Sample, ToyLLM, apply_bug, apply_op, bar,  # noqa: E402
                       extend_trace, fmt_size, majority_vote, make_problem, make_problem_set, mean, pad,
                       sample_step)

CHECK_T = 1.0     # PRM 打分时用的温度：1.0 = checker 原本的概率（校准的"似然"）
RESOLVE_T = 0.3   # ORM 重解时用的温度：仔细地做


def step_prob(checker, step, x, written, step_index, temperature=CHECK_T):
    """
    checker 从输入 x 出发，写出 written 这个值的概率。这就是 PRM 给这一步的分数。

    【思路】三种做法各自会得到什么值是已知的：correct → 正确值，bug → 错误套路的值，
        careless → 正确值 + 6 种偏移之一（每种概率相同）。把"得到 written"的做法的概率加起来即可：
            P(written) = P(correct)·[written == 正确值] + P(bug)·[written == bug 值]
                         + P(careless) × (偏移中能得到 written 的个数) / 6
    【参数】checker：ToyLLM；step；x：这一步的输入（候选自己写下的上一步的值）；written：候选写下的这一步的值；
        step_index；temperature
    【返回】0~1 的 float
    【边界】
        · 这是**局部**检查：只问"在这个输入上，这一步算得对不对"，不管输入本身对不对。
          前面错了、这一步在错误的输入上算对了，分数仍然高 —— 所以要把各步的分数合起来看（aggregate）。
        · checker 的盲点会传给 PRM：如果 checker"最相信"错误套路，它会给错误的一步打高分。
    """
    probs = checker.step_probs(step.op, x, step.operand, step_index, temperature)
    right = apply_op(step.op, x, step.operand)
    p = 0.0
    if written == right:
        p += probs[0]
    if written == apply_bug(step.op, x, step.operand):
        p += probs[1]
    p += probs[2] * sum(1 for d in CARELESS_OFFSETS if right + d == written) / len(CARELESS_OFFSETS)
    return p


def prm_step_scores(checker, problem, sample, temperature=CHECK_T):
    """
    PRM：给候选的每一步打分（step_prob），返回分数列表（和步数等长）。

    【参数】checker；problem；sample：候选（Sample）；temperature
    【返回】list of float
    """
    scores = []
    x = problem.start
    for i, step in enumerate(problem.steps):
        scores.append(step_prob(checker, step, x, sample.values[i], i, temperature))
        x = sample.values[i]
    return scores


def aggregate(step_scores, how="product"):
    """
    把各步的分数合成一个解答的分数。

    【参数】step_scores：各步分数；how："product"（乘积，Lightman et al. 2023）/ "min"（最差的一步）/
        "last"（最后一步，Snell et al. 2024 的一种选择）
    【返回】float
    【思路】乘积 = "每一步都对"的概率（如果各步的判断独立）；min 只看最可疑的一步；
        两者都会惩罚"任何一步可疑"的解答。last 只看最后一步，等于把 PRM 当 ORM 用。
    【边界】乘积对长解答不公平：步数越多，乘积越小（即使每一步都很可信）。比较不同长度的解答时要注意。
    """
    if how == "product":
        return math.prod(step_scores)
    if how == "min":
        return min(step_scores)
    if how == "last":
        return step_scores[-1]
    raise ValueError(f"未知的聚合方式：{how!r}")


def orm_score(checker, problem, sample, rng, m=4, temperature=RESOLVE_T):
    """
    ORM：checker 自己从头把题做 m 遍，返回"得到和候选相同答案"的比例。

    【参数】checker；problem；sample；rng；m：重解次数；temperature：重解时的温度
    【返回】0~1 的 float（只有 m + 1 种可能的值：0, 1/m, ..., 1）
    【边界】
        · 它只看最终答案，看不到过程，也无法指出哪一步错了。
        · 要判断 k 步题，checker 自己也要连对 k 步；checker 不比生成者强时，它对难题几乎没有判断力。
        · 每次打分要花 m 次完整的生成，比 PRM（k 次单步检查）贵。
    """
    return sum(checker.generate(problem, rng, temperature).answer == sample.answer for _ in range(m)) / m


def first_error(problem, sample):
    """
    候选里第一个算错的步骤的下标（局部检查，Sample.local_ok）；全对返回 None。

    【思路】Lightman et al. 2023 的人工标注就是这样：标注员从第一步开始逐步判断，标到第一个错误为止。
    """
    for i, ok in enumerate(sample.local_ok):
        if not ok:
            return i
    return None


def on_track(problem, sample):
    """
    每一步之后，解答是否还"在正确的路上"（到这一步为止的值都等于正确值）。返回 list of bool。

    【思路】这是"过程标签"最自然的定义：第 i 步的前缀如果全对，从这里继续还有可能做对；
        一旦偏离（某一步错了），后面就几乎不可能回到正确答案（除非碰巧）。
    """
    good, out = True, []
    for v, truth in zip(sample.values, problem.correct_values()):
        good = good and (v == truth)
        out.append(good)
    return out


def rollout_value(completer, problem, sample, i, rng, n, temperature=0.7):
    """
    Math-Shepherd 式的 rollout：保留候选的前 i + 1 步，让 completer 续写 n 次，返回做对的比例。

    【参数】completer：续写用的模型；problem；sample；i：保留到第几步（0 ≤ i < 步数）；rng；n：续写次数
    【返回】0~1 的 float。
        soft estimation（SE）= 这个比例本身；hard estimation（HE）= 这个比例 > 0（至少一次做对）。
    【边界】
        · i 是最后一步时，不需要续写：直接看答案对不对（n 次结果都一样）。
        · 这个值估计的是"从这里出发，completer 还能做对的概率"，即 completer 的**价值函数**，
          而不是"这一步本身对不对"。前缀全对、但剩下的步骤很难时，SE 也会很低。
    """
    if i == problem.n_steps - 1:
        return 1.0 if sample.correct else 0.0
    prefix_v, prefix_p = list(sample.values[:i + 1]), list(sample.procs[:i + 1])
    hits = 0
    for _ in range(n):
        if extend_trace(completer, problem, prefix_v, prefix_p, rng, temperature).correct:
            hits += 1
    return hits / n


def select_by_scores(samples, scores, weighted=True):
    """
    按分数挑选候选。weighted=True 时同一答案的分数相加（weighted best-of-N），否则选单个最高分。

    【参数】samples：Sample 列表；scores：等长的分数列表（越高越好，最好在 0~1 之间）
    【返回】被选中的 Sample
    """
    if not weighted:
        return samples[max(range(len(samples)), key=lambda i: scores[i])]
    totals, first = {}, {}
    for s, sc in zip(samples, scores):
        totals[s.answer] = totals.get(s.answer, 0.0) + sc
        first.setdefault(s.answer, s)
    return first[max(totals, key=lambda a: totals[a])]


if __name__ == "__main__":
    import random
    rng = random.Random(0)
    problem = make_problem(rng, 5)
    gen = ToyLLM(1e10)
    print(problem.text(), "标准答案：", problem.answer(), " 正确的中间值：", problem.correct_values())
    for _ in range(4):
        s = gen.generate(problem, rng, 0.7)
        print(f"\n候选：{list(s.values)}  {'✓' if s.correct else '✗'}  做法 {', '.join(s.procs)}")
        for name, checker in (("PRM（10B 自查）", gen), ("PRM（100B）", ToyLLM(1e11))):
            sc = prm_step_scores(checker, problem, s)
            print(f"  {pad(name, 16)} 各步 {[round(x, 2) for x in sc]}  乘积 {aggregate(sc):.3f}")
        print(f"  {pad('ORM（10B 重解 4 次）', 16)} {orm_score(gen, problem, s, rng):.2f}")