"""
=============================================================================
CS486 · 第 4 讲 · 代码 03：搜索与算力分配——同样的预算，怎样花最划算？
=============================================================================

【这个脚本要回答的问题】（讲义 §6–§7）
    测试时计算不只有"采样 N 次再挑"一种花法。两个问题：
    1. **搜索**：与其一次生成整条推理链再打分，不如**每走一步就让验证器看一眼**，只保留有希望的几条
       （逐步束搜索；Snell et al. 2024、Wu et al. 2024 的 PRM 引导搜索）。在同样的计算量下，它比 best-of-N 好吗？
    2. **分配**：一批题有难有易。与其每道题都采样同样多次，不如先试一两次，把预算留给"还没把握"的题
       （Snell et al. 2024 的"计算最优"分配）。能省多少？

    演示 1：贪心、强制 k 步的贪心（利用结构知识的对照）、三种打分方式的束搜索、best-of-N——准确率和计算量。
    演示 1b：验证器在"策略模型很少写出的候选"上可靠吗？
    演示 2：按"样本是否一致"自适应地分配采样预算 vs 平均分配 vs 按真实难度分配（看不到的上限）。

【运行】
    python3 03_search_and_budget.py           # 约 1 分钟（需要代码 02 训练好的验证器；没有会先训练约 2 分钟）
    python3 03_search_and_budget.py --quick   # 冒烟测试

【依赖】PyTorch + 本目录的 ttc_lab.py + 第 1、2 讲的模块。
"""

import argparse
import os
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

import hop_world as hw                                                   # noqa: E402
import tiny_reasoner as tr                                               # noqa: E402
import ttc_lab as tl                                                     # noqa: E402
from console import pct, print_header, print_observations, print_table  # noqa: E402

KS = (3, 4, 5, 6)


def bon_prm_full(probs, samples, scores, n):
    """用前 n 个样本、按 PRM-full 的乘积分数挑最高的，返回准确率。"""
    ok = 0
    for p, outs, sc in zip(probs, samples, scores):
        j = max(range(n), key=lambda q: tl.prm_aggregate(sc[q]))
        ok += tl.is_correct(p, outs[j])
    return ok / len(probs)


def demo1(policy, vers, probs, samples, scores, quick):
    print_header(f"演示 1：搜索 vs 采样（n = 12，k ∈ {KS}，{len(probs)} 道题）")
    res = tr.evaluate(policy, probs, "cot")
    greedy = tr.accuracy(res)
    greedy_tokens = sum(len(r["generated"].split(hw.END)[0]) + 1 for r in res)
    forced = tl.greedy_forced_length(policy, probs)
    forced_acc = sum(a == p.answer for a, p in zip(forced, probs)) / len(probs)
    rows = [["贪心", pct(greedy), f"{greedy_tokens / len(probs):.1f}", "0"],
            ["贪心，强制走 k 步就停（结构知识，不是测试时计算）", pct(forced_acc),
             f"{sum(p.k for p in probs) / len(probs):.1f}", "0"]]
    out = {"greedy": greedy, "forced": forced_acc}
    for label, beam, expand, wp, wv in (("束搜索 4×4：只用策略概率", 4, 4, 1.0, 0.0),
                                        ("束搜索 4×4：只用 PRM-full", 4, 4, 0.0, 1.0),
                                        ("束搜索 2×2：策略概率 + PRM-full", 2, 2, 1.0, 1.0),
                                        ("束搜索 4×4：策略概率 + PRM-full", 4, 4, 1.0, 1.0)):
        ans, cost = tl.beam_search(policy, vers["prm_full"], probs, beam=beam, expand=expand,
                                   w_policy=wp, w_verifier=wv)
        acc = sum(a == p.answer for a, p in zip(ans, probs)) / len(probs)
        out[label] = (acc, cost)
        rows.append([label, pct(acc), f"{cost['policy_tokens'] / len(probs):.1f}",
                     f"{cost['verifier_calls'] / len(probs):.1f}"])
    N = len(samples[0])
    for n in (4, 16, N):
        if n > N:
            continue
        acc = bon_prm_full(probs, samples, scores, n)
        toks = sum(len(o.split(hw.END)[0]) + 1 for outs in samples for o in outs[:n]) / len(probs)
        out[("bon", n)] = (acc, toks)
        rows.append([f"best-of-{n}（PRM-full 选最高）", pct(acc), f"{toks:.1f}", f"{n}"])
    print_table(["方法", "准确率", "每题策略 token", "每题验证器打分次数"], rows, [52, 8, 14, 18], ["<", ">", ">", ">"])
    print("（策略 token：策略模型一共做了多少次'生成一个 token'的前向；束搜索每一轮对每条束各做一次。"
          "\n  束搜索 b×e：束宽 b，每条束取 e 个候选；'|'（停下）也是候选之一。）")
    return out


def first_step_probe(vers, quick):
    """
    PRM-full 能不能分辨"正确的第一步"和"错误的第一步"？

    【思路】对每道题，构造只有一步的推理链：一条写对了第 1 跳，三条写成了别的字母。
        策略模型自己几乎从不在第一步犯错，所以验证器的训练数据里很少有"错误的第一步"——这是一个**离策略**（off-policy）的探针。
    【返回】dict：(n, k) → (正确候选的分数中位数, 错误候选的分数中位数)
    """
    print_header("演示 1b：PRM-full 在'策略模型很少写出的候选'上的打分（只写一步的推理链）")
    rows, out = [], {}
    for n in (8, 12):
        probs = hw.make_eval_set(n, 4, 10 if quick else 100, seed=21)
        good, bad = [], []
        for p in probs:
            wrong = [c for c in hw.LETTERS if c != p.walk[1]][:3]
            exs = [tl.verifier_example(p, p.walk[1] + "|" + p.walk[1] + "$")] + \
                  [tl.verifier_example(p, w + "|" + w + "$") for w in wrong]
            sc = tl.score(vers["prm_full"], exs, "prm_full")
            good.append(sc[0][0])
            bad += [x[0] for x in sc[1:]]
        med_g, med_b = sorted(good)[len(good) // 2], sorted(bad)[len(bad) // 2]
        out[n] = (med_g, med_b)
        rows.append([f"n = {n}，k = 4", f"{med_g:.3f}", f"{med_b:.3f}"])
    print_table(["题目", "正确第一步的分数（中位数）", "错误第一步的分数（中位数）"], rows, [16, 26, 26])
    return out


def allocate(probs, answers, budget, pilot=4, oracle_rates=None):
    """
    在一批题之间分配采样预算，最后用多数投票作答。

    【策略】
        oracle_rates 为 None（"按一致性自适应"）：每道题先采 pilot 个样本；pilot 个答案完全一致的题就此停下，
            剩下的预算平均分给答案不一致的题——"大家意见一致"被当作"这题简单、已经会了"的信号。
        oracle_rates 给定（"按真实难度分配"，一个看不到的上限）：知道每道题的真实成功率，
            只把额外预算给"有时对、有时错"（0 < 成功率 < 1）的题，"总是对"和"从来不对"的题只用 pilot 个。
    【参数】answers：每道题的样本答案列表（样本池）；budget：平均每道题的样本数
    【返回】(准确率, 实际平均用量, 拿到额外预算的题数)
    【边界】每道题最多用到样本池的大小；没有题需要额外预算时，剩余预算不再使用。
    """
    N = len(answers[0])
    used = [min(pilot, N)] * len(probs)
    if oracle_rates is None:
        extra_ids = [i for i, a in enumerate(answers) if len(set(a[:pilot])) > 1]
    else:
        extra_ids = [i for i, r in enumerate(oracle_rates) if 0 < r < 1]
    remaining = budget * len(probs) - sum(used)
    if extra_ids and remaining > 0:
        each = remaining // len(extra_ids)
        for j, i in enumerate(extra_ids):
            used[i] = min(N, used[i] + each + (1 if j < remaining % len(extra_ids) else 0))
    acc = sum(tl.majority_vote(answers[i][:used[i]]) == probs[i].answer for i in range(len(probs))) / len(probs)
    return acc, sum(used) / len(probs), len(extra_ids)


def demo2(policy, probs, quick):
    print_header("演示 2：在一批题之间分配采样预算（温度 1.5，多数投票作答；先给每道题 4 个样本）")
    N = 4 if quick else 32
    samples = tl.sample_many(policy, probs, N, 1.5, seed=102)          # 和代码 01、02 的温度 1.5 样本相同
    answers = [[hw.parse_answer(o, "cot") for o in outs] for outs in samples]
    rates = [sum(tl.is_correct(p, o) for o in outs) / N for p, outs in zip(probs, samples)]
    rows, out = [], {}
    for budget in (4, 8, 16):
        if budget > N:                       # 【边界】预算超过样本池的大小（--quick）时跳过
            continue
        uni = sum(tl.majority_vote(a[:budget]) == p.answer for a, p in zip(answers, probs)) / len(probs)
        ada, used_a, n_a = allocate(probs, answers, budget)
        ora, used_o, n_o = allocate(probs, answers, budget, oracle_rates=rates)
        out[budget] = (uni, ada, ora, n_a, n_o)
        rows.append([budget, pct(uni), f"{pct(ada)}（{n_a} 题加预算）", f"{pct(ora)}（{n_o} 题加预算）"])
    print_table(["平均每题预算", "平均分配", "按一致性自适应", "按真实难度分配（上限）"], rows, [14, 10, 24, 26])
    hopeless = sum(1 for i, a in enumerate(answers) if len(set(a[:4])) > 1 and rates[i] == 0)
    n_dis = sum(1 for a in answers if len(set(a[:4])) > 1)
    print(f"（前 4 个样本意见不一致的 {n_dis} 道题里，有 {hopeless} 道 32 次采样一次都没对过。）")
    return out, hopeless, n_dis


def main():
    ap = argparse.ArgumentParser(description="搜索与算力分配")
    ap.add_argument("--quick", action="store_true", help="冒烟测试")
    args = ap.parse_args()
    t0 = time.perf_counter()
    per_k, N = (5, 4) if args.quick else (50, 32)

    print_header("准备：策略模型、PRM-full 验证器、测试题与样本池（温度 1.0，和代码 01、02 相同）")
    policy = tr.get_model("cot_L2", quick=args.quick)[0]
    vers = tl.load_verifiers(quick=args.quick)
    probs = []
    for k in KS:
        probs += hw.make_eval_set(12, k, per_k, seed=21)
    samples = tl.sample_many(policy, probs, N, 1.0, seed=101)
    flat = [tl.verifier_example(p, o) for p, outs in zip(probs, samples) for o in outs]
    s = tl.score(vers["prm_full"], flat, "prm_full")
    scores = [s[i * N:(i + 1) * N] for i in range(len(probs))]

    d1 = demo1(policy, vers, probs, samples, scores, args.quick)
    probe = first_step_probe(vers, args.quick)
    d2, hopeless, n_dis = demo2(policy, probs, args.quick)

    if args.quick:
        print("\n（--quick 冒烟测试：数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    bp, cp = d1["束搜索 4×4：只用策略概率"]
    bv, _ = d1["束搜索 4×4：只用 PRM-full"]
    bc, cc = d1["束搜索 4×4：策略概率 + PRM-full"]
    bc2, cc2 = d1["束搜索 2×2：策略概率 + PRM-full"]
    bon32, tok32 = d1[("bon", N)]
    print_observations([
        f"贪心 {pct(d1['greedy'])}；只要强制它走 k 步就停，贪心就到了 {pct(d1['forced'])}——代码 01 发现的'走过头'失败，"
        "用一条结构知识就修好了大部分。这不是测试时计算，但提醒我们：先弄清楚模型错在哪里，有时一个简单的约束胜过大量计算。",
        f"束搜索的打分方式决定成败：只用策略概率 {pct(bp)}，只用 PRM-full **{pct(bv)}**（比贪心还差得多），"
        f"两者相加 {pct(bc)}。只看验证器时，搜索会把策略模型几乎不会写出的字母也拿来比较，而 PRM-full 是在策略模型自己的样本上"
        f"训练的，从没见过这类候选——它在那里是'瞎'的（演示 1b：n = 12 时错误第一步的分数中位数 {probe[12][1]:.3f}，"
        f"正确的是 {probe[12][0]:.3f}；n = 8 时分别是 {probe[8][1]:.3f} 和 {probe[8][0]:.3f}）。"
        "加上策略概率，候选就被限制在'模型认为合理'的范围里，验证器只负责在其中挑。",
        f"只用策略概率的束搜索也比贪心好很多（{pct(bp)} vs {pct(d1['greedy'])}）。一个可能的原因：束搜索的分数是各步对数概率之和，"
        "每多走一步就多加一个负数，所以它偏爱更短的完整答案，这恰好抵消了模型'走过头'的倾向——如果是这样，"
        "这是一个'歪打正着'的效果，不能推广（这里没有单独检验，见讲义练习 5）。",
        f"计算量：策略概率 + PRM-full 的束搜索（4×4）每题约 {cc['policy_tokens'] / len(probs):.0f} 次策略前向、"
        f"{cc['verifier_calls'] / len(probs):.0f} 次验证器打分，准确率 {pct(bc)}；best-of-32 用约 {tok32:.0f} 个策略 token、"
        f"32 次打分，只有 {pct(bon32)}。2×2 的小束搜索（{cc2['policy_tokens'] / len(probs):.0f} 次前向）也有 {pct(bc2)}。"
        "把验证器的信号用在每一步，比只在最后挑一次有效得多——前提是搜索不把验证器带出它见过的分布。",
        f"预算分配（演示 2）：按'前 4 个样本是否一致'自适应地分配，没有比平均分配更好（平均每题 8 个样本时分别是 "
        f"{pct(d2[8][1])} 和 {pct(d2[8][0])}）。原因是：意见不一致的 {n_dis} 道题里有 {hopeless} 道其实从来没对过——"
        "'不一致'既可能表示'难但有希望'，也可能表示'没希望'，额外的预算有很大一部分浪费在后者身上。"
        f"即使知道真实难度、只给'有时对'的题加预算，平均每题 8 / 16 个样本也只到 {pct(d2[8][2])} / {pct(d2[16][2])}"
        f"（平均分配 {pct(d2[8][0])} / {pct(d2[16][0])}）：多数投票本身限制了能从额外样本里挑出多少。"
        "计算最优分配的价值取决于**难度估计**有多准（Snell et al. 2024 也发现测试时计算主要帮助'中等难度'的题）。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()