"""
=============================================================================
第 4 讲 · 代码 04：self-consistency —— 多想几遍、少数服从多数
=============================================================================

对应讲义：lectures/lecture04_chain_of_thought.md §4.3

【这个脚本要回答的问题】
    Wang et al.（2023）的 self-consistency：同一个问题，用温度采样生成多条思维链，取最终答案的多数票。
    它为什么有效？什么时候无效？在 scratch 思维链模型上：
    1. 采样 1 / 3 / 5 / 9 / 15 条，多数投票的准确率怎样变化？和贪心解码比呢？（演示 1）
    2. 投票能修正哪些错误、修正不了哪些？按"正确答案是不是这道题最常见的答案"把题目分类。（演示 2）
    3. 多数投票的数学：每次采样答对的概率是 p，错误答案有多分散——投票之后的准确率是多少？（演示 3）

【怎么读输出】
    三组题：训练过的组合、没训练过的组合（都是用首都 / 语言做桥）、领导人做桥（模型不知道答案，反转诅咒）。
    "投票 @m"：取前 m 条采样的最终答案，出现次数最多的那个（并列时取先出现的）；没有给出答案的采样不参与投票。

【运行】python3 04_self_consistency.py     （读 02 号脚本训练好的缓存；约 15 秒）
【依赖】PyTorch；import 同目录的 cot_lab.py，第 1 讲的 console.py。
"""

import math
import os
import sys
from collections import Counter

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
from console import print_header, print_observations, print_table  # noqa: E402
from cot_lab import (COMBOS, EVAL_COUNTRIES, TEST_COMBOS, TRAIN_COMBOS, ask, attr, get_cot_model,  # noqa: E402
                     parse_answer, render_chat, twohop_question)

M_MAX = 15
MS = [1, 3, 5, 9, 15]
GROUPS = {
    "训练过的组合": [x for x in TRAIN_COMBOS if x[0] != "leader"],
    "没训练过的组合": [x for x in TEST_COMBOS if x[0] != "leader"],
    "领导人做桥": [x for x in COMBOS if x[0] == "leader"],
}


def vote(answers):
    """多数投票：忽略 None；出现次数最多的答案，并列时取最先出现的那个。全是 None 时返回 None。"""
    xs = [a for a in answers if a is not None]
    if not xs:
        return None
    counts = Counter(xs)
    best = max(counts.values())
    return next(a for a in xs if counts[a] == best)


def collect(m, tok, temperature):
    """每道题：贪心答案 + M_MAX 条采样的答案。返回 {组名: [(标准答案, 贪心答案, [采样答案...]), ...]}。"""
    out = {}
    for g, combos in GROUPS.items():
        rows = []
        for b, t in combos:
            for c in EVAL_COUNTRIES:
                p = render_chat(twohop_question(b, t, c))
                greedy = parse_answer(ask(m, tok, p), "scratch")
                samples = [parse_answer(ask(m, tok, p, temperature=temperature, seed=s), "scratch") for s in range(M_MAX)]
                rows.append((attr(c, t), greedy, samples))
        out[g] = rows
    return out


def demo1(data_by_temp):
    """演示 1：投票曲线。"""
    print_header("演示 1：采样 m 条思维链、多数投票（测试国家；每组 32 道题）")
    rows, res = [], {}
    for temp, data in data_by_temp.items():
        for g, items in data.items():
            greedy = sum(gr == gold for gold, gr, _ in items) / len(items)
            single = sum(sum(s == gold for s in ss) for gold, _, ss in items) / (len(items) * M_MAX)
            accs = [sum(vote(ss[:m]) == gold for gold, _, ss in items) / len(items) for m in MS]
            res[(temp, g)] = (greedy, single, accs)
            rows.append([f"T = {temp}", g, f"{100 * greedy:.1f}%", f"{100 * single:.1f}%"] + [f"{100 * a:.1f}%" for a in accs])
    print_table(["温度", "题目", "贪心", "单条采样（平均）"] + [f"投票 @{m}" for m in MS], rows)
    u7 = res[(0.7, "没训练过的组合")]
    u1 = res[(1.0, "没训练过的组合")]
    l1 = res[(1.0, "领导人做桥")]
    print_observations([
        f"没训练过的组合（T = 1.0）：单条采样平均 {100 * u1[1]:.1f}%，比贪心的 {100 * u1[0]:.1f}% 低——采样会抽到不太可能的中间国家。"
        f"投票 @15 回到 {100 * u1[2][-1]:.1f}%：多数投票把分散的错误投掉了大部分，但**没有超过贪心解码**。",
        "为什么没超过？在 scratch 格式里，每个答案只对应一条'路径'（写下哪个国家，就答哪个国家的属性）。贪心解码挑的是"
        "概率最大的那个国家，也就是概率最大的那个答案；投票只是用 15 次采样去估计同一件事，还带着噪声。",
        "Wang et al.（2023）的关键洞察在于：真实的推理问题里，**很多条不同的推理路径会通向同一个正确答案**（同一道算术题有好几种解法），"
        "而错误的路径各错各的。概率最大的那一条**路径**（贪心）未必通向概率最大的那个**答案**；对路径求和（投票）才是在找最可能的答案。"
        "05 号脚本的模拟器演示这种情形。",
        f"温度 0.7 时单条采样更接近贪心（{100 * u7[1]:.1f}%），投票的收益也更小。领导人做桥：贪心 {100 * l1[0]:.1f}%，"
        f"投票 @15 是 {100 * l1[2][-1]:.1f}%——模型不知道这些国家，采样只是在不同的错误之间换来换去，**投票修不好'不知道'**。",
        "Wang et al. 报告 self-consistency 在思维链的基础上又带来明显提升（例如 PaLM 540B 在 GSM8K 上提高了约 18 个百分点），"
        "代价是 m 倍的调用费用。",
    ])
    return res


def demo2(data):
    """演示 2：哪些题投票修得好、哪些修不好。"""
    print_header("演示 2：按 15 条采样里'最常见的答案是不是正确答案'把题目分类（T = 1.0）")
    rows = []
    share_ok, share_bad = [], []
    for g, items in data.items():
        plur_right = plur_wrong = 0
        fixed = broken = 0
        for gold, greedy, ss in items:
            counts = Counter(a for a in ss if a is not None)
            top = counts.most_common(1)[0][0] if counts else None
            if top == gold:
                plur_right += 1
            else:
                plur_wrong += 1
            v = vote(ss)
            share = counts.most_common(1)[0][1] / M_MAX if counts else 0.0
            (share_ok if v == gold else share_bad).append(share)
            fixed += (greedy != gold) and (v == gold)
            broken += (greedy == gold) and (v != gold)
        rows.append([g, plur_right, plur_wrong, fixed, broken])
    print_table(["题目", "正确答案是最常见的答案", "某个错误答案最常见", "贪心错 → 投票对", "贪心对 → 投票错"], rows)
    tot_fixed, tot_broken = sum(r[3] for r in rows), sum(r[4] for r in rows)
    print_observations([
        f"投票相对贪心：修好了 {tot_fixed} 道、弄坏了 {tot_broken} 道（三组合计）。在这个'一个答案一条路径'的模型上，投票基本是在复现贪心。",
        "投票最终会收敛到每道题'最常见的答案'。所以 m 足够大时，投票的准确率 = '正确答案是最常见答案'的题目比例：",
        "对模型'大体知道'的题（训练过、没训练过的组合），正确答案通常是最常见的，投票（相对单条采样）把偶然的错误修掉；",
        "对模型'不知道'的题（领导人做桥），总有某个错误答案更常见（它反复编出同一个国家），投票反而把这个错误固定下来——"
        "这就是**系统性错误**：每次犯同一个错，再多的采样也投不掉。",
        f"投票结果的'一致程度'（最常见答案占 15 票里的多少）本身是一个信号：投票答对的题平均 {100 * sum(share_ok) / len(share_ok):.0f}%，"
        f"答错的题平均 {100 * sum(share_bad) / len(share_bad):.0f}%。一致程度低的题更可能答错——第 5 讲的 SelfCheckGPT 就用这个思路检测幻觉。",
    ])


def p_majority(p, k_wrong, m):
    """
    多数投票答对的概率（精确计算）：每次采样以概率 p 答对，以 (1 − p)/k_wrong 的概率答成 k_wrong 个错误答案中的某一个。

    【思路】枚举正确答案得到的票数 j；在剩下的 m − j 票里，要求每个错误答案的票数都 < j
        （并列时算答错，偏保守）。用动态规划数"把 m − j 票分给 k_wrong 个错误答案、且每个都 < j"的方式。
    【边界】k_wrong = 1 时就是普通的二项分布多数票。m 较大、k_wrong 较大时计算量是 O(m² · k_wrong)，这里足够快。
    """
    q = (1 - p) / k_wrong if k_wrong else 0.0
    total = 0.0
    for j in range(1, m + 1):
        rest = m - j
        # ways[r] = 把 r 票分给已经考虑过的错误答案（每个 < j 票）的"概率权重"（多项分布的系数 × q^r）
        ways = [1.0] + [0.0] * rest
        for _ in range(k_wrong):
            new = [0.0] * (rest + 1)
            for r in range(rest + 1):
                if ways[r] == 0.0:
                    continue
                for x in range(min(j - 1, rest - r) + 1):
                    new[r + x] += ways[r] * (q ** x) / math.factorial(x)
            ways = new
        total += (p ** j) / math.factorial(j) * ways[rest] * math.factorial(m)
    return total


def demo3():
    """演示 3：多数投票的数学。"""
    print_header("演示 3：多数投票的数学——单次答对的概率 p、错误答案的个数 k，投票之后的准确率")
    rows = []
    for p, k in [(0.6, 1), (0.6, 5), (0.4, 5), (0.4, 1), (0.3, 10), (0.2, 10)]:
        rows.append([f"{p:.1f}", k, f"{(1 - p) / k:.2f}"] + [f"{100 * p_majority(p, k, m):.1f}%" for m in [1, 5, 15, 41]])
    print_table(["单次答对 p", "错误答案的个数 k", "每个错误答案的概率", "投票 @1", "@5", "@15", "@41"], rows)
    print_observations([
        "p = 0.6、错误集中在 1 个答案上：投票让准确率随 m 慢慢上升（@41 时 90.3%）——这是孔多塞陪审团定理：只要每个'投票者'"
        "比随机好、彼此独立，人越多越准。错误分散在 5 个答案上时，上升得快得多（@15 时 98.8%）。",
        "p = 0.4、但错误分散在 5 个答案上（每个只有 0.12）：正确答案仍是最常见的，投票照样能把准确率推高。"
        "**投票要求的不是 p > 0.5，而是'正确答案比任何一个错误答案都更常见'。**",
        "p = 0.4、错误集中在 1 个答案上（0.6）：错误答案更常见，投票越多越错，最终趋近 0——这就是演示 2 里领导人做桥的那种题。",
        "成本：投票 @m 要 m 次生成（每次都要完整的思维链）。所以实际系统常用小的 m（5–10），或者先用一次生成、只在不确定时再多采样"
        "（自适应的 self-consistency）。",
    ])


def main():
    m, tok, _ = get_cot_model(fmt="scratch")
    data = {temp: collect(m, tok, temp) for temp in (0.7, 1.0)}
    demo1(data)
    demo2(data[1.0])
    demo3()


if __name__ == "__main__":
    main()