"""
=============================================================================
CS486 · 第 7 讲 · 代码 03：模型"知道自己不知道"吗？——置信度、校准与选择性回答
=============================================================================

【这个脚本要回答的问题】（讲义 §5）
    元认知（metacognition）：模型对自己的答案有多大把握，这个把握准不准？
    Kadavath et al.（2022，"Language Models (Mostly) Know What They Know"）发现大模型在多选题上的概率**校准**得不错，
    还能预测"我能不能答对这道题"（P(IK)）。但"有区分度"和"校准"是两回事，而且一个看起来有区分度的信号，
    可能只是和某个表面特征（例如长度）相关。

    演示 1：区分度——各种置信度信号区分"对 / 错"的 AUC；以及在"回答长度相同"的题里还剩多少区分度。
    演示 2：校准——答案字母的概率，和实际准确率对得上吗？分布内（n = 8）vs 分布外（n = 12、16）。
    演示 3：选择性回答——只回答最有把握的一部分题，准确率能提高多少？
    演示 4：无解的题（第 6 讲代码 03）——模型说不出"无解"，但它的概率里有没有"不对劲"的信号？

【运行】
    python3 03_know_what_you_know.py           # 约 15 秒
    python3 03_know_what_you_know.py --quick   # 冒烟测试

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

import argparse
import os
import sys
import time
from collections import defaultdict

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

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

SIGNALS = ("seq_logp", "mean_logp", "min_prob", "ans_prob", "neg_len", "orm", "prm")
LABEL = {"seq_logp": "整条回答的对数概率", "mean_logp": "平均每 token 对数概率", "min_prob": "最小 token 概率",
         "ans_prob": "答案字母的概率", "neg_len": "回答越短（对照）", "orm": "ORM（第 4 讲）", "prm": "PRM-full（第 4 讲）"}


def greedy_resps(model, probs):
    return [rl.response_of(o) for o in tr.generate(model, [p.prompt for p in probs], max(p.k for p in probs) + 5)]


def stratified_auc(scores, labels, groups):
    """
    分层 AUC：只在同一组（这里是"回答长度相同"）内部比较对的和错的，再按每组的"对 × 错"配对数加权平均。

    【思路】如果一个信号的区分度全部来自"错的回答更长"，那么在长度相同的回答之间它就分不出对错，分层 AUC 会接近 0.5。
    【返回】(分层 AUC, 参与比较的配对数)；没有任何组同时有对有错时返回 (None, 0)
    """
    by = defaultdict(list)
    for s, y, g in zip(scores, labels, groups):
        by[g].append((s, y))
    num, den = 0.0, 0
    for items in by.values():
        pos = [s for s, y in items if y]
        neg = [s for s, y in items if not y]
        if pos and neg:
            pairs = len(pos) * len(neg)
            num += sc.auc([s for s, _ in items], [y for _, y in items]) * pairs
            den += pairs
    return (num / den, den) if den else (None, 0)


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 = 10、12、14（k = 3–6，各 %d 道）" % per_k)
    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(10, per_k=per_k) + rl.eval_set(12, per_k=per_k) + rl.eval_set(14, per_k=per_k)
    resps = greedy_resps(policy, probs)
    ok = [sc.is_right(p, r) for p, r in zip(probs, resps)]
    sig = sc.confidence_signals(policy, probs, resps, vers)
    print(f"  贪心准确率 {pct(sum(ok) / len(ok))}（{len(probs)} 道题）")

    print_header("演示 1：区分度——AUC（全部题目 vs 只在回答长度相同的题之间比较）")
    lens = [len(r) for r in resps]
    rows, aucs = [], {}
    for s in SIGNALS:
        a = sc.auc(sig[s], ok)
        st, pairs = stratified_auc(sig[s], ok, lens)
        aucs[s] = (a, st)
        rows.append([LABEL[s], f"{a:.3f}", "-" if st is None else f"{st:.3f}"])
    print_table(["信号", "AUC", "同长度内的 AUC"], rows, [22, 8, 16])
    _, pairs = stratified_auc(sig["seq_logp"], ok, lens)
    n_pairs = sum(ok) * (len(ok) - sum(ok))
    print(f"  （同长度内的 AUC 只用了 {pairs} 个'对–错'配对，占全部 {n_pairs} 个配对的 {pct(pairs / max(1, n_pairs))}："
          "大多数对的和错的回答长度不同）")

    print_header("演示 2：校准——答案字母的概率 vs 实际准确率（ECE：期望校准误差，越小越好）")
    rows, cal = [], {}
    for n in (8, 12, 16):
        ps = rl.eval_set(n, per_k=per_k)
        rs = greedy_resps(policy, ps)
        y = [sc.is_right(p, r) for p, r in zip(ps, rs)]
        conf = sc.confidence_signals(policy, ps, rs)["ans_prob"]
        e, table = sc.ece(conf, y)
        cal[n] = (e, sum(conf) / len(conf), sum(y) / len(y), table)
        rows.append([f"n = {n}", f"{sum(conf) / len(conf):.3f}", pct(sum(y) / len(y)), f"{e:.3f}"])
    print_table(["测试集", "平均置信度", "实际准确率", "ECE"], rows, [10, 12, 12, 8])
    e16, _, _, table16 = cal[16]
    print("  n = 16 的可靠性表（置信度箱子：样本数、平均置信度、准确率）：")
    for lo, cnt, mc, ma in table16:
        print(f"    [{lo:.1f}, {lo + 0.1:.1f})：{cnt:>3} 道，平均置信度 {mc:.3f}，准确率 {pct(ma)}")

    print_header("演示 3：选择性回答——只回答最有把握的那部分题（演示 1 的 %d 道题）" % len(probs))
    covs = (0.25, 0.5, 0.75, 1.0)
    rows, selc = [], {}
    for s in ("seq_logp", "neg_len", "orm"):
        r = sc.selective(sig[s], ok, covs)
        selc[s] = r
        rows.append([LABEL[s]] + [pct(r[c]) for c in covs])
    print_table(["按什么排序", *[f"回答 {pct(c, 0)}" for c in covs]], rows, [22, 12, 12, 12, 12])

    print_header("演示 4：无解的题——模型的概率里有没有'不对劲'的信号？（第 6 讲 distract，n = 16、k = 6）")
    model6 = rlab.sft_model("distract", quick=args.quick)
    items = rlab.sweep_problems(6, 40 if args.quick else 400, drop_from=5, drop_prob=0.5, seed=77)
    ps = [p for p, _ in items]
    solvable = [s for _, s in items]
    rs = greedy_resps(model6, ps)
    s6 = sc.confidence_signals(model6, ps, rs)
    rows, unsolv = [], {}
    for s in ("min_prob", "seq_logp", "ans_prob"):
        a = sc.auc(s6[s], solvable)
        unsolv[s] = a
        mean_ok = sum(v for v, y in zip(s6[s], solvable) if y) / max(1, sum(solvable))
        mean_bad = sum(v for v, y in zip(s6[s], solvable) if not y) / max(1, len(solvable) - sum(solvable))
        rows.append([LABEL[s], f"{mean_ok:.3f}", f"{mean_bad:.3f}", f"{a:.3f}"])
    print_table(["信号", "可解的题（平均）", "无解的题（平均）", "AUC（可解 vs 无解）"], rows, [22, 18, 18, 20])
    # 两种阈值：一个"直觉"的 0.5；一个在**开发集**（另一个种子生成的扫描题）上按 Youden 指数选的阈值
    dev_items = rlab.sweep_problems(6, 40 if args.quick else 400, drop_from=5, drop_prob=0.5, seed=78)
    dev_rs = greedy_resps(model6, [p for p, _ in dev_items])
    dev_min = sc.confidence_signals(model6, [p for p, _ in dev_items], dev_rs)["min_prob"]
    dev_sol = [ok_ for _, ok_ in dev_items]
    best_j, th_dev = -1.0, 0.5
    for t in sorted(set(dev_min)):
        c = sum(v < t and not y for v, y in zip(dev_min, dev_sol)) / max(1, len(dev_sol) - sum(dev_sol))
        f = sum(v < t and y for v, y in zip(dev_min, dev_sol)) / max(1, sum(dev_sol))
        if c - f > best_j:
            best_j, th_dev = c - f, t
    for name, th in (("直觉阈值", 0.5), ("开发集上选的阈值", th_dev)):
        flagged = [v < th for v in s6["min_prob"]]
        caught = sum(f and not y for f, y in zip(flagged, solvable)) / max(1, len(solvable) - sum(solvable))
        false = sum(f and y for f, y in zip(flagged, solvable)) / max(1, sum(solvable))
        unsolv[name] = (th, caught, false)
        print(f"  规则'最小 token 概率 < {th:.3f} 就说无解'（{name}）：标出 {pct(caught)} 的无解题，误标 {pct(false)} 的可解题")
    print(f"  （测试集共 {len(solvable)} 道，其中 {len(solvable) - sum(solvable)} 道无解；开发集用另一个种子生成）")
    ex = next((p, r, s6["min_prob"][i]) for i, (p, r) in enumerate(zip(ps, rs)) if not solvable[i])
    missing = next(c for c in ex[0].walk[:-1] if c not in {st[0] for st in ex[0].statements})
    print(f"  例：题面缺了 {missing!r} 的出边；模型写 {ex[1]!r}，最小 token 概率 {ex[2]:.3f}")

    if args.quick:
        print("\n（--quick 冒烟测试：数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    th_d, c_d, f_d = unsolv["开发集上选的阈值"]
    _, c_n, _ = unsolv["直觉阈值"]
    answered16 = [row for row in table16 if row[0] >= 0.9]
    print_observations([
        f"模型自己的置信度看上去很有区分度（演示 1）：整条回答的对数概率 AUC {aucs['seq_logp'][0]:.3f}——可是'回答越短越可能对'这个"
        f"和模型无关的规则就有 {aucs['neg_len'][0]:.3f}。只在回答长度相同的题之间比较时，它降到 {aucs['seq_logp'][1]:.3f}，"
        f"答案字母的概率降到 {aucs['ans_prob'][1]:.3f}（接近随机）；第 4 讲的验证器仍有 {aucs['orm'][1]:.3f}（ORM）/ "
        f"{aucs['prm'][1]:.3f}（PRM-full）。**模型的'自知'大部分来自长度**：走过头的回答更长，长回答的总概率自然更低。"
        "评估一个置信度信号时，要和简单的表面特征比较，并在控制这些特征之后再看。",
        f"分布内校准很好，分布外严重过度自信（演示 2）：n = 8 时平均置信度 {cal[8][1]:.3f}、准确率 {pct(cal[8][2])}；"
        f"n = 16 时 200 道题里只有 {sum(r[1] for r in answered16)} 道给出了置信度 ≥ 0.9 的答案（其余大多一直没停下、没给答案），"
        f"这些答案的平均置信度 {answered16[0][2]:.3f}，准确率却只有 {pct(answered16[0][3])}。系统性的错误往往伴随着很高的把握——"
        "第 4 讲的投票、本讲的二值纠错都对它无效，原因相同。",
        f"选择性回答（演示 3）：只回答自己最有把握的 25% 的题，准确率从 {pct(selc['seq_logp'][1.0])} 升到 {pct(selc['seq_logp'][0.25])}"
        f"（ORM 排序 {pct(selc['orm'][0.25])}，只按长度排序 {pct(selc['neg_len'][0.25])}）。'知道什么时候不回答'是有用的，"
        "但在这里它和'挑短的回答'几乎一样好——实际部署时，这样的弃答规则在长度分布变化时会失效。",
        f"无解的题（演示 4）：第 6 讲的 distract 说不出'无解'，会编造一步；但它的回答没那么有把握了：回答里最小的 token 概率"
        f"平均从可解题的 {sum(v for v, y in zip(s6['min_prob'], solvable) if y) / max(1, sum(solvable)):.3f} 降到 "
        f"{sum(v for v, y in zip(s6['min_prob'], solvable) if not y) / max(1, len(solvable) - sum(solvable)):.3f}，区分可解 / 无解的 AUC {unsolv['min_prob']:.3f}。用'< 0.5 就说无解'这条直觉规则只能标出 "
        f"{pct(c_n)}；在开发集上选阈值（{th_d:.3f}）后，标出 {pct(c_d)} 的无解题、误标 {pct(f_d)} 的可解题。"
        "**模型'察觉到不对劲'的信号是存在的，只是落在一个没有校准过的尺度上**——要用它，先校准（第 2 讲：阈值在开发集上选）。",
        "元认知的三个层次：察觉（我可能错了）、定位（错在哪一步）、校准（我有多大把握）。这个小模型在察觉上有一些真实的信号"
        "（同长度内仍有约 0.70 的 AUC、无解题上的概率下降），在定位上几乎没有（代码 02），在分布外完全没有校准。"
        "Kadavath et al.（2022）在大模型上看到的'大致知道自己知道什么'，正是在这三个层次上需要分别检验的。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()