"""
=============================================================================
CS486 · 第 7 讲 · 代码 02：反馈从哪里来、说了什么——在实验台上做自我纠错
=============================================================================

【这个脚本要回答的问题】（讲义 §3–§4）
    cot_L2 在 n = 12 的长题面上贪心只有 29.5%，主要错在"走过头"（第 4 讲）。给它一轮又一轮的"纠错"机会，能改对多少？
    我们比较两类反馈、三种来源：

        反馈的内容：
            二值（binary）：只说"这个答案可以 / 不可以接受"；不可以就重新采样一次（温度 1.0）
            定位（localized）：说"第 j 步错了"；在那一步换成策略的下一个候选 token，再接着写
        反馈的来源：
            自己（self）  ：策略模型自己的概率（整条回答的对数概率 / 每个 token 的概率）
            验证器        ：第 4 讲的 ORM（二值）和 PRM-full（定位）
            神谕（oracle）：求解器——知道真实答案和每一步对不对（上界，真实使用中没有）

    演示 1：二值反馈 + 重新采样，6 轮。
    演示 2：定位反馈 + 换一个候选，6 轮。
    演示 3：第一轮的四个参数（代码 01 的模型）：检查的 TPR / FPR，修复率 / 改坏率——瓶颈在哪一个？
    演示 4：对照——不纠错，而是利用结构知识（强制走 k 步）或训练（第 5 讲的 GRPO 模型）。

【运行】
    python3 02_feedback_sources.py           # 约 30 秒
    python3 02_feedback_sources.py --quick   # 冒烟测试

【依赖】PyTorch + 本目录的 selfcorrect_lab.py；第 4 讲的验证器缓存、第 5 讲的 GRPO 模型缓存（没有时会先训练）。
"""

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 selfcorrect_lab as sc                                             # noqa: E402  先导入它：它把第 2、4、5 讲的目录加进 sys.path
import hop_world as hw                                                   # noqa: E402
import rl_lab as rl                                                      # 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

ROUNDS = 6


def youden_threshold(scores, labels):
    """
    在开发集上选一个阈值：使 TPR − FPR（Youden 指数）最大，"分数 ≥ 阈值"判为可以接受。

    【为什么用开发集】阈值在测试集上选，就是在测试集上调参（第 2 讲）。开发集和测试集用不同的种子生成。
    【返回】阈值（float）
    """
    best, best_t = -1.0, None
    pos = sum(labels)
    neg = len(labels) - pos
    for t in sorted(set(scores)):
        tpr = sum(1 for s, y in zip(scores, labels) if y and s >= t) / max(1, pos)     # 对的被接受
        fpr = sum(1 for s, y in zip(scores, labels) if not y and s >= t) / max(1, neg)  # 错的被接受
        if tpr - fpr > best:
            best, best_t = tpr - fpr, t
    return best_t


def history_table(title, results):
    print_header(title)
    rows = []
    for name, (_, hist) in results.items():
        cells = [name]
        by_round = {h[0]: h for h in hist}
        last = hist[-1] if hist else None
        for r in (1, 2, 4, 6):
            h = by_round.get(r, last)
            cells.append("-" if h is None else f"{pct(h[1], 1)}（+{h[2]} / −{h[3]}）")
        rows.append(cells)
    print_table(["反馈来源", "1 轮", "2 轮", "4 轮", "6 轮"], rows, [20, 18, 18, 18, 18])
    print("  （每格：准确率（错→对 的题数 / 对→错 的题数），都相对于纠错之前的贪心回答；提前结束的沿用最后一轮）")


def round_one_params(probs, before, after, flagged):
    """
    从第一轮的实际结果估计代码 01 的四个参数。

    【返回】dict：tpr（错的回答被标出的比例）、fpr（对的回答被标出的比例）、
                 p_fix（被标出的错回答变对的比例）、p_break（被标出的对回答变错的比例）、
                 pred（用这四个数算出的一轮后准确率）、actual（实际的一轮后准确率）、acc0
    """
    ok0 = [sc.is_right(p, r) for p, r in zip(probs, before)]
    ok1 = [sc.is_right(p, r) for p, r in zip(probs, after)]
    wrong = [i for i, o in enumerate(ok0) if not o]
    right = [i for i, o in enumerate(ok0) if o]
    fw = [i for i in wrong if flagged[i]]
    fr = [i for i in right if flagged[i]]
    tpr = len(fw) / max(1, len(wrong))
    fpr = len(fr) / max(1, len(right))
    p_fix = sum(ok1[i] for i in fw) / max(1, len(fw))
    p_break = sum(not ok1[i] for i in fr) / max(1, len(fr))
    acc0 = sum(ok0) / len(ok0)
    pred, _ = sc.one_round_gain(acc0, tpr, fpr, p_fix, p_break)
    return {"tpr": tpr, "fpr": fpr, "p_fix": p_fix, "p_break": p_break, "pred": pred,
            "actual": sum(ok1) / len(ok1), "acc0": acc0}


def main():
    ap = argparse.ArgumentParser(description="自我纠错：反馈的来源与内容")
    ap.add_argument("--quick", action="store_true", help="冒烟测试")
    args = ap.parse_args()
    t0 = time.perf_counter()
    per_k = 5 if args.quick else 50

    print_header("准备：cot_L2、第 4 讲的验证器；n = 12 的测试集（和第 4、5 讲相同）与开发集（不同的种子）")
    policy = tr.get_model("cot_L2", quick=args.quick, verbose=not args.quick)[0]
    vers = tl.load_verifiers(quick=args.quick, verbose=not args.quick)
    probs = rl.eval_set(12, per_k=per_k)
    dev = []
    for k in (3, 4, 5, 6):
        dev += hw.make_eval_set(12, k, per_k, seed=22)
    base = [rl.response_of(o) for o in tr.generate(policy, [p.prompt for p in probs], 11)]
    dev_resps = [rl.response_of(o) for o in tr.generate(policy, [p.prompt for p in dev], 11)]
    acc0 = sum(sc.is_right(p, r) for p, r in zip(probs, base)) / len(probs)
    dev_sig = sc.confidence_signals(policy, dev, dev_resps)
    dev_ok = [sc.is_right(p, r) for p, r in zip(dev, dev_resps)]
    tau = youden_threshold(dev_sig["seq_logp"], dev_ok)
    print(f"  贪心准确率：测试集 {pct(acc0)}，开发集 {pct(sum(dev_ok) / len(dev_ok))}；"
          f"'自己'的二值检查：整条回答的对数概率 ≥ {tau:.3f} 就接受（在开发集上按 Youden 指数选的阈值）")

    def self_accept(p, r):
        return sc.confidence_signals(policy, [p], [r])["seq_logp"][0] >= tau

    def orm_accept(p, r):
        return tl.score(vers["orm"], [tl.verifier_example(p, r)], "orm")[0] >= 0.5

    binary = {
        "自己（对数概率）": sc.resample_retry(policy, probs, base, self_accept, ROUNDS),
        "ORM（≥ 0.5）": sc.resample_retry(policy, probs, base, orm_accept, ROUNDS),
        "神谕（答案对吗）": sc.resample_retry(policy, probs, base, sc.is_right, ROUNDS),
    }
    history_table(f"演示 1：二值反馈 + 重新采样（{len(probs)} 道题；纠错前 {pct(acc0)}）", binary)

    local = {
        "自己（token 概率 < 0.5）": sc.localized_correction(policy, probs, base, "self", rounds=ROUNDS),
        "PRM-full（< 0.5）": sc.localized_correction(policy, probs, base, "prm", vers, rounds=ROUNDS),
        "神谕（第一个错步）": sc.localized_correction(policy, probs, base, "oracle", rounds=ROUNDS),
    }
    history_table("演示 2：定位反馈 + 在那一步换一个候选", local)

    print_header("演示 3：第一轮的四个参数（代码 01 的模型）——瓶颈在哪一个？")
    rows, params = [], {}
    first_binary = {
        "二值·自己": [not self_accept(p, r) for p, r in zip(probs, base)],
        "二值·ORM": [not orm_accept(p, r) for p, r in zip(probs, base)],
        "二值·神谕": [not sc.is_right(p, r) for p, r in zip(probs, base)],
    }
    after1 = {
        "二值·自己": sc.resample_retry(policy, probs, base, self_accept, 1)[0],
        "二值·ORM": sc.resample_retry(policy, probs, base, orm_accept, 1)[0],
        "二值·神谕": sc.resample_retry(policy, probs, base, sc.is_right, 1)[0],
    }
    for name, kind in (("定位·自己", "self"), ("定位·PRM-full", "prm"), ("定位·神谕", "oracle")):
        scores = sc.step_scores(kind, policy, vers, probs, base)
        first_binary[name] = [any(v < 0.5 for v in s) for s in scores]
        after1[name] = sc.localized_correction(policy, probs, base, kind, vers, rounds=1)[0]
    for name in first_binary:
        q = round_one_params(probs, base, after1[name], first_binary[name])
        params[name] = q
        rows.append([name, f"{q['tpr']:.2f}", f"{q['fpr']:.2f}", f"{q['p_fix']:.2f}", f"{q['p_break']:.2f}",
                     pct(q["pred"]), pct(q["actual"])])
    print_table(["方法", "TPR", "FPR", "修复率", "改坏率", "按公式复原", "实际"], rows, [16, 6, 6, 8, 8, 10, 8])
    print("  （对一轮来说，代码 01 的公式是一个精确的分解：用实测的四个参数代回去，必然等于实际准确率。它的用处是看出瓶颈在哪个因子）")
    print("  （TPR：错的回答被标出的比例；FPR：对的回答被标出的比例；修复率 / 改坏率：被标出的错 / 对回答，一轮后变对 / 变错的比例）")

    print_header("演示 4：对照——不在测试时纠错，而是利用结构知识或训练")
    forced = tl.greedy_forced_length(policy, probs)
    acc_forced = sum(a == p.answer for a, p in zip(forced, probs)) / len(probs)
    grpo_job = next(j for j in rl.main_jobs() if j["method"] == "grpo" and j["seed"] == 0)
    grpo = rl.load_runs([grpo_job], quick=args.quick, verbose=False)[0][0]
    acc_grpo = rl.greedy_accuracy(grpo, probs)
    best_local = local["神谕（第一个错步）"][1][-1][1]
    print_table(["做法", "准确率"], [
        ["贪心（不纠错）", pct(acc0)],
        ["最好的测试时纠错（神谕定位，6 轮）", pct(best_local)],
        ["强制在第 k 步停下（第 4 讲：结构知识）", pct(acc_forced)],
        ["第 5 讲 GRPO 训练后的模型（种子 0），贪心", pct(acc_grpo)],
    ], [40, 10])

    if args.quick:
        print("\n（--quick 冒烟测试：数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    b_self, b_orc = params["二值·自己"], params["二值·神谕"]
    l_self, l_prm, l_orc = params["定位·自己"], params["定位·PRM-full"], params["定位·神谕"]
    print_observations([
        f"二值反馈 + 重新采样几乎完全无效（演示 1）：6 轮之后，'自己' {pct(binary['自己（对数概率）'][1][-1][1])}、"
        f"ORM {pct(binary['ORM（≥ 0.5）'][1][-1][1])}，连神谕（准确地知道哪些答案错了）也只有 {pct(binary['神谕（答案对吗）'][1][-1][1])}。"
        f"瓶颈不在检查（演示 3：'自己'的检查 TPR {b_self['tpr']:.2f}，神谕 {b_orc['tpr']:.2f}），而在修复：被标出的错误回答，"
        f"重新采样之后只有 {b_orc['p_fix']:.2f} 变对。cot_L2 在这些题上的错误是**系统性**的（走过头），温度 1.0 下重新采样还是同一个错——"
        "第 4 讲的 pass@k 曲线几乎是平的，原因相同。**知道'错了'而不知道'错在哪、该怎么改'，再试一次没有用。**",
        f"定位反馈有效，但取决于定位得准不准（演示 2）：神谕告诉模型第一个错步，6 轮后 {pct(local['神谕（第一个错步）'][1][-1][1])}"
        f"（{local['神谕（第一个错步）'][1][-1][2]} 道错→对，0 道对→错）；PRM-full 定位，6 轮后 {pct(local['PRM-full（< 0.5）'][1][-1][1])}，"
        f"但它误标了 {pct(l_prm['fpr'], 0)} 的正确回答，这些回答**全部**被改坏（改坏率 {l_prm['p_break']:.2f}）——"
        "在一条正确的推理链上'换一个候选'，必然把它换错。这对应 Tyen et al.（2024）的发现：模型找不到自己的推理错误，"
        "但告诉它错在哪里，它往往能改对；以及代码 01 的结论：改坏率高时，检查的 FPR 决定成败。",
        f"模型自己的概率能**察觉**错误，却不能**定位**错误：整条回答的对数概率作为二值检查，TPR {b_self['tpr']:.2f}、FPR {b_self['fpr']:.2f}；"
        f"可是逐个 token 看，只有 {pct(l_self['tpr'], 0)} 的错误回答里有一个概率低于 0.5 的 token——模型在每一步都很'自信'，"
        "包括走过头的那一步。整条回答的概率之所以能区分对错，主要是因为错的回答更长（代码 03 演示 1 检验这一点）。",
        f"所有测试时的纠错都不如利用结构知识或训练（演示 4）：最好的纠错（神谕定位，6 轮）{pct(best_local)}，"
        f"强制在第 k 步停下 {pct(acc_forced)}，第 5 讲 GRPO 训练后的模型 {pct(acc_grpo)}。"
        "这和 SCoRe（Kumar et al. 2024）等工作的出发点一致：与其指望模型在测试时纠正自己，不如**训练**它一开始就做对，"
        "或者训练它学会纠错（而不是靠提示）。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()