"""
=============================================================================
CS486 · 第 4 讲 · 代码 01：多采样几次——覆盖率、温度与多数投票
=============================================================================

【这个脚本要回答的问题】（讲义 §2–§3）
    同一个模型，同一道题，回答 N 次（温度 > 0），能不能比回答 1 次更好？
    两个不同的问题要分开：
        覆盖率（coverage，pass@N）：N 次里**至少有一次**对的概率——"答案在不在样本里"；
        选择（selection）：从 N 个样本里**挑出**对的那个——不知道答案时，只能靠投票或验证器。
    Brown et al.（2024，Large Language Monkeys）发现覆盖率随采样次数近似"对数线性"地增长；
    Wang et al.（2023，自洽性）发现多数投票能稳定地提高准确率。在我们的实验台上看看这些现象。

    策略模型：第 1 讲的 cot_L2。题目：n = 12 的长题面（它不完美的地方），k = 3、4、5、6 各 50 道。

    演示 1：温度的影响——贪心 vs 温度 0.5 / 1.0 / 1.5：单次准确率、pass@8、pass@32、多数投票@32。
    演示 2：覆盖率和投票随 N 的变化（温度 1.0），按跳数分开看。
    演示 3：每道题的"成功率"分布——为什么投票有时有用、有时没用？
    演示 4：投票什么时候会出错？
    演示 5：失败的类型——错在"查找"，还是错在"何时停下"？

【运行】
    python3 01_sampling_and_coverage.py           # 约 40 秒（采样约 2 万条推理链）
    python3 01_sampling_and_coverage.py --quick   # 冒烟测试

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

import argparse
import os
import sys
import time
from collections import Counter

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 test_problems(per_k):
    """n = 12、k = 3–6 各 per_k 道（固定种子）。"""
    out = []
    for k in KS:
        out += hw.make_eval_set(12, k, per_k, seed=21)
    return out


def summarize(probs, samples, ns=(1, 8, 32)):
    """
    对每道题的样本算 pass@n（无偏估计）和多数投票@N。

    【返回】dict：{"pass": {n: 平均 pass@n}, "maj": 多数投票@全部样本 的准确率, "rates": 每道题的成功率}
    """
    N = len(samples[0])
    rates, passes, maj = [], {n: 0.0 for n in ns if n <= N}, 0
    for p, outs in zip(probs, samples):
        c = sum(tl.is_correct(p, o) for o in outs)
        rates.append(c / N)
        for n in passes:
            passes[n] += tl.pass_at_k(N, c, n)
        maj += tl.majority_vote([hw.parse_answer(o, "cot") for o in outs]) == p.answer
    m = len(probs)
    return {"pass": {n: v / m for n, v in passes.items()}, "maj": maj / m, "rates": rates}


def demo1(policy, probs, N):
    print_header(f"演示 1：温度的影响（{len(probs)} 道题，每道采样 {N} 次）")
    greedy = tr.evaluate(policy, probs, "cot")
    g_acc = tr.accuracy(greedy)
    rows = [["贪心（温度 0）", pct(g_acc), "-", "-", "-"]]
    out = {"greedy": g_acc, "greedy_res": greedy}
    for i, T in enumerate((0.5, 1.0, 1.5)):
        samples = tl.sample_many(policy, probs, N, T, seed=100 + i)
        s = summarize(probs, samples, ns=(1, 8, N))
        out[T] = (s, samples)
        rows.append([f"温度 {T}", pct(s["pass"][1]), pct(s["pass"].get(8)), pct(s["pass"][N]), pct(s["maj"])])
    print_table(["解码方式", "单次准确率（pass@1）", "pass@8", f"pass@{N}", f"多数投票@{N}"], rows,
                [16, 22, 10, 10, 14])
    return out


def demo2(probs, samples, N):
    print_header(f"演示 2：覆盖率与投票随采样次数 N 的变化（温度 1.0），按跳数分开")
    rows, curves = [], {}
    for k in KS:
        idx = [i for i, p in enumerate(probs) if p.k == k]
        for n in (1, 2, 4, 8, 16, 32):
            if n > N:
                continue
            cov = sum(tl.pass_at_k(N, sum(tl.is_correct(probs[i], o) for o in samples[i]), n) for i in idx) / len(idx)
            maj = sum(tl.majority_vote([hw.parse_answer(o, "cot") for o in samples[i][:n]]) == probs[i].answer
                      for i in idx) / len(idx)
            curves[(k, n)] = (cov, maj)
    for n in (1, 2, 4, 8, 16, 32):
        if n > N:
            continue
        rows.append([n] + [f"{pct(curves[(k, n)][0], 0)} / {pct(curves[(k, n)][1], 0)}" for k in KS])
    print_table(["N"] + [f"k = {k}：覆盖 / 投票" for k in KS], rows, [6, 22, 22, 22, 22])
    print("（覆盖 = pass@N 的无偏估计；投票 = 前 N 个样本的多数投票是否正确。）")
    return curves


def demo3(rates):
    print_header("演示 3：每道题的成功率（温度 1.0 时 32 个样本里对的比例）的分布")
    bins = [("0（从没对过）", lambda r: r == 0), ("(0, 0.25]", lambda r: 0 < r <= 0.25),
            ("(0.25, 0.5]", lambda r: 0.25 < r <= 0.5), ("(0.5, 0.75]", lambda r: 0.5 < r <= 0.75),
            ("(0.75, 1)", lambda r: 0.75 < r < 1), ("1（每次都对）", lambda r: r == 1)]
    rows, out = [], {}
    for name, f in bins:
        c = sum(1 for r in rates if f(r))
        out[name] = c / len(rates)
        rows.append([name, c, pct(c / len(rates)), "#" * round(40 * c / len(rates))])
    print_table(["成功率", "题数", "比例", ""], rows, [16, 6, 8, 42], ["<", ">", ">", "<"])
    return out


def demo4(probs, samples, greedy_res):
    print_header("演示 4：投票什么时候出错？（温度 1.0，32 个样本）")
    cats = Counter()
    wrong_plurality = []
    for p, outs, g in zip(probs, samples, greedy_res):
        answers = [hw.parse_answer(o, "cot") for o in outs]
        c = sum(a == p.answer for a in answers) / len(answers)
        vote = tl.majority_vote(answers)
        if c == 0:
            cats["从没对过（投票不可能对）"] += 1
        elif vote == p.answer:
            cats["正确答案得票最多"] += 1
            if c < 0.5:
                cats["  其中：正确答案不到半数，仍然最多"] += 1
        else:
            cats["有对的样本，但某个错误答案得票更多"] += 1
            top = Counter(a for a in answers if a).most_common(1)[0][1] / len(answers)
            wrong_plurality.append((c, top))
        cats["贪心对"] += g["correct"]
    rows = [[k, v, pct(v / len(probs))] for k, v in cats.items()]
    print_table(["情况", "题数", "比例"], rows, [38, 6, 8], ["<", ">", ">"])
    if wrong_plurality:
        avg_c = sum(c for c, _ in wrong_plurality) / len(wrong_plurality)
        avg_t = sum(t for _, t in wrong_plurality) / len(wrong_plurality)
        print(f"（投票选错的题里：正确答案的平均得票率 {pct(avg_c)}，得票最多的错误答案平均 {pct(avg_t)}。）")
    return cats, wrong_plurality


def failure_type(p, output):
    """
    把一条样本归入五类之一（讲义 §3.3）：
        正确；
        每一步查找都对，但步数太多（没有及时停下：k 步之后还在走，或者一直没写 "|"）；
        每一步查找都对，但步数太少；
        推理链完全正确（恰好 k 步），但 "|" 后面的答案抄错了；
        至少有一步查找错了。
    【思路】这个世界里，"查找"和"何时停下（数到 k）"是两种不同的能力。只看对错分不出来，
        要看推理链本身：长度对不对、每一步对不对。
    """
    if tl.is_correct(p, output):
        return "正确"
    trace = tl.trace_of(output)
    if all(tl.step_labels(p, trace)):
        if len(trace) > p.k:
            return "查找都对，但走过头了"
        return "查找都对，但停早了" if len(trace) < p.k else "推理链对，答案抄错了"
    return "有一步查找错了"


def demo5(probs, samples):
    print_header("演示 5：失败的类型（温度 1.0 的全部样本）——错在'查找'还是错在'何时停下'？")
    cats = ("正确", "查找都对，但走过头了", "查找都对，但停早了", "推理链对，答案抄错了", "有一步查找错了")
    rows, out = [], {}
    for k in KS:
        cnt = Counter()
        for p, outs in zip(probs, samples):
            if p.k == k:
                for o in outs:
                    cnt[failure_type(p, o)] += 1
        total = sum(cnt.values())
        out[k] = {c: cnt[c] / total for c in cats}
        rows.append([f"k = {k}"] + [pct(cnt[c] / total) for c in cats])
    print_table(["跳数"] + list(cats), rows, [8, 8, 22, 20, 20, 16])
    return out


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)
    policy = tr.get_model("cot_L2", quick=args.quick)[0]
    probs = test_problems(per_k)
    print_header("准备")
    print(f"  策略模型 cot_L2（第 1 讲；训练时 n ≤ 8）。测试题：n = 12，k ∈ {KS}，各 {per_k} 道。")
    print(f"  例：{probs[0].prompt}  答案 {probs[0].answer}")

    d1 = demo1(policy, probs, N)
    s10, samples10 = d1[1.0]
    curves = demo2(probs, samples10, N)
    hist = demo3(s10["rates"])
    cats, wp = demo4(probs, samples10, d1["greedy_res"])
    fails = demo5(probs, samples10)

    if args.quick:
        print("\n（--quick 冒烟测试：数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    s05, s15 = d1[0.5][0], d1[1.5][0]
    n = len(probs)
    print_observations([
        f"贪心 {pct(d1['greedy'])}；温度 1.0 时单次采样 {pct(s10['pass'][1])}，但 32 次里至少有一次对的概率（pass@32）是 "
        f"{pct(s10['pass'][32])}，多数投票@32 是 {pct(s10['maj'])}。在这个温度下多采样只把覆盖率提高了几个点，"
        "投票又把其中大部分'挑'了出来——覆盖率（答案在不在样本里）和选择（能不能挑出来）是两件事。",
        f"温度：0.5 时 pass@32 只有 {pct(s05['pass'][32])}（样本太像，覆盖率低），1.5 时单次准确率降到 {pct(s15['pass'][1])}"
        f"（样本太乱），pass@32 {pct(s15['pass'][32])}。多样性和质量之间有权衡；测试时计算需要'既敢于不同、又不至于胡来'的温度。",
        f"按跳数看（演示 2）：k = 3 时单次就有 {pct(curves[(3, 1)][0], 0)}，k = 5 时覆盖率从 {pct(curves[(5, 1)][0], 0)} 涨到 "
        f"{pct(curves[(5, 32)][0], 0)}；奇怪的是 k = 4 和 k = 6 几乎是 0（{pct(curves[(4, 32)][0], 0)}、{pct(curves[(6, 32)][0], 0)}），"
        "比更'难'的 k = 5 还差。演示 5 解释了原因。",
        f"失败的类型（演示 5）：k = 4 时 {pct(fails[4]['查找都对，但走过头了'])}、k = 6 时 {pct(fails[6]['查找都对，但走过头了'])} 的样本"
        "**每一步查找都是对的**，只是没有在第 k 步停下——在比训练时长的题面上，模型数不清自己走了几步。"
        f"其余的失败是至少有一步查找错了（k = 4：{pct(fails[4]['有一步查找错了'])}；k = 6：{pct(fails[6]['有一步查找错了'])}）。"
        "这种错误是**系统性**的（模型很确定地继续往下走），所以多采样、投票都救不了它；"
        "而一个只检查'每一步查找对不对'的过程验证器也看不出它（代码 02）。",
        f"成功率的分布是两极化的（演示 3）：{pct(hist['0（从没对过）'])} 的题 32 次一次都没对，{pct(hist['1（每次都对）'])} 的题每次都对，"
        "中间的题才是测试时计算真正起作用的地方。对'从没对过'的题，再多采样、再好的选择方法都没用——"
        "测试时计算放大的是模型已有的能力，不能凭空创造能力。",
        f"投票选错的情况（演示 4）：{cats['有对的样本，但某个错误答案得票更多']} 道题里有对的样本，却有一个错误答案得票更多；"
        f"另一方面，有 {cats['  其中：正确答案不到半数，仍然最多']} 道题正确答案不到半数仍然胜出——"
        "只要错误答案足够分散。投票依赖的不是'正确率 > 50%'，而是'正确答案比任何一个错误答案都多'。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()