"""
=============================================================================
第 12 讲 · 01：不用提示的 CoT —— 在第一个 token 处分支，按"答案的置信度"挑选
=============================================================================

对应讲义：lectures/lecture12_llm_reasoning.md 第 3 节（Wang & Zhou 2024, Chain-of-Thought Reasoning Without Prompting）

-----------------------------------------------------------------------------
【设定：一个没有被提示"一步一步想"的预训练模型】
-----------------------------------------------------------------------------
题目是第 1 讲的多步算术题。模型的**第一个 token** 决定了它接下来怎样回答：
    · 直接回答（direct）：第一个 token 就是答案的开头，后面不写推理（第 1 讲的 direct 模式，
      一次前向传播能可靠完成的步数有限，步数一多就容易错）；
    · 先推理（CoT）：第一个 token 是推理的开头（"我有 3 个苹果，爸爸比我多 2 个……"），
      后面一步一步写出来（第 1 讲的 CoT 模式）。
第一个 token 的候选按概率从高到低排成 K = 10 个分支。**概率最高的那个（贪心解码会选它）总是直接回答**；
其余每个分支以 P_COT_BRANCH 的概率是推理的开头 —— 这是论文的关键观察：
"推理路径其实就在模型的解码空间里，只是贪心解码没有走到它们"。
分支之后继续解码（temperature = 0.7，和前几讲一样：toy world 里的贪心解码不现实地强）。

-----------------------------------------------------------------------------
【答案的置信度 Δ】
-----------------------------------------------------------------------------
论文：Δ = 答案 token 上"第 1 名和第 2 名的概率差"的平均。toy 版本：
    · CoT 路径：答案是最后一步算出来的值。在写最后一步时，模型对"这一步结果"的分布是：
      正确结果 p_correct、固定错误套路的结果 p_bug、六种手滑的结果各 p_careless / 6。
      Δ = P(选中的结果) − P(其他结果里最大的那个)。
    · 直接回答：答案是"在脑子里"一次算完的。模型对答案的分布用 64 次采样来估计（temperature = 1），
      Δ = 选中的答案的频率 − 其他答案里最大的频率。
      多步运算全压在一次前向传播里，答案的分布很分散，Δ 通常很小。

【比较的做法】（200 道 3~6 步的题，三种模型规模）
    贪心（不提示）                ：走概率最高的分支 —— 直接回答
    CoT-decoding（最大 Δ）        ：K 个分支都走一遍，选 Δ 最大的那条的答案
    CoT-decoding（按答案聚合 Δ）  ：把答案相同的路径的 Δ 加起来，选总和最大的答案
    自洽性（不提示，10 个样本）    ：按第一个 token 的概率随机选分支，走 10 次，多数投票
    CoT 提示（1 个样本）           ：用 few-shot CoT 提示，每次都先推理
    CoT 提示 + 自洽性（10 个样本） ：第 2 讲的 self-consistency
    K 个分支的覆盖率（oracle）     ：K 条路径里只要有一条对就算对（挑选的上限）

运行时间：约 3 秒（只用标准库）。
"""

import math
import os
import random
import sys

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

from toy_llm import (CARELESS_OFFSETS, ToyLLM, apply_bug, apply_op, fmt_size, majority_vote, make_problem_set,  # noqa: E402
                     mean, pad, sample_index)

K_BRANCH = 10               # 第一个 token 的候选数（论文里的 top-k，k = 10）
FIRST_TOKEN_DECAY = 0.45    # 第 i 个候选的概率 ∝ exp(−0.45 i)：排名越靠后，概率越小
P_COT_BRANCH = 0.5          # 排名第 2~10 的候选，每个以 0.5 的概率是"推理的开头"
TEMPERATURE = 0.7           # 分支之后继续解码的温度
N_PROBE = 64                # 估计"直接回答"的答案分布用的样本数
N_SC = 10                   # 自洽性用的样本数
N_PROBLEMS = 200


def first_token_probs():
    """
    第一个 token 的 K 个候选的概率：p_i ∝ exp(−FIRST_TOKEN_DECAY × i)，i = 0, 1, …, K−1。

    【返回】长度 K 的 list，和为 1，从大到小排列。第 0 个是贪心解码会选的那个。
    """
    weights = [math.exp(-FIRST_TOKEN_DECAY * i) for i in range(K_BRANCH)]
    total = math.fsum(weights)
    return [w / total for w in weights]


def branch_types(rng):
    """
    一道题的 K 个分支各是什么类型。

    【返回】长度 K 的 list，元素是 "direct" 或 "cot"；第 0 个总是 "direct"。
    【思路】论文发现：对数学题，贪心解码（第 1 名的 token）大多直接给出答案；推理路径常常出现在
        排名靠后的候选里。哪些分支是推理的开头，每道题不一样，这里随机决定。
    """
    return ["direct"] + ["cot" if rng.random() < P_COT_BRANCH else "direct" for _ in range(K_BRANCH - 1)]


def last_step_distribution(model, problem, values):
    """
    CoT 路径在写**最后一步**时，模型对这一步结果的分布。

    【参数】model：CoT 模式的 ToyLLM；problem：题目；values：这条路径每一步写出来的值
    【返回】dict {结果: 概率}。正确结果、bug 的结果、六种手滑的结果；相同的结果概率相加。
    【思路】最后一步的输入是**这条路径自己写出来的**倒数第二个值（可能已经错了）。
        模型在这个输入上"正确地"算，就得到 right；所以即使前面错了，最后一步的置信度仍可能很高 ——
        置信度只能反映"答案 token 本身"有多确定，不能发现前面步骤里的错误。
    【边界】只有一步的题，最后一步的输入就是起始数。
    """
    i = problem.n_steps - 1
    step = problem.steps[i]
    x = values[-2] if problem.n_steps > 1 else problem.start
    p_correct, p_bug, p_careless = model.step_probs(step.op, x, step.operand, i, temperature=1.0)
    right = apply_op(step.op, x, step.operand)
    dist = {right: p_correct}
    bug = apply_bug(step.op, x, step.operand)
    dist[bug] = dist.get(bug, 0.0) + p_bug
    for off in CARELESS_OFFSETS:
        dist[right + off] = dist.get(right + off, 0.0) + p_careless / len(CARELESS_OFFSETS)
    return dist


def margin(dist, chosen):
    """
    Δ = P(选中的答案) − P(其他答案里最大的)。

    【参数】dist：{答案: 概率或频率}；chosen：路径实际给出的答案
    【返回】float，介于 −1 和 1 之间。选中的答案不是分布里的第 1 名时为负数。
    【边界】分布里只有选中的答案一个时，"其他答案里最大的"按 0 算。
    """
    others = [p for a, p in dist.items() if a != chosen]
    return dist.get(chosen, 0.0) - (max(others) if others else 0.0)


def direct_distribution(model_direct, problem, rng):
    """
    用 N_PROBE 次采样估计"直接回答"时模型对答案的分布（temperature = 1，即模型本身的分布）。

    【返回】dict {答案: 频率}
    【思路】真实的 LLM 可以直接读出答案 token 的概率；toy 模型没有"答案 token"，只能用采样来估计。
    """
    counts = {}
    for _ in range(N_PROBE):
        a = model_direct.generate(problem, rng, temperature=1.0).answer
        counts[a] = counts.get(a, 0) + 1
    return {a: c / N_PROBE for a, c in counts.items()}


def decode_branch(kind, model_cot, model_direct, problem, probe, rng):
    """
    沿一个分支解码到底。

    【参数】kind："cot" 或 "direct"；probe：direct_distribution 的结果（同一道题的直接回答共用）
    【返回】(答案, 是否正确, Δ)
    """
    if kind == "cot":
        s = model_cot.generate(problem, rng, temperature=TEMPERATURE)
        return s.answer, s.correct, margin(last_step_distribution(model_cot, problem, s.values), s.answer)
    s = model_direct.generate(problem, rng, temperature=TEMPERATURE)
    return s.answer, s.correct, margin(probe, s.answer)


def run_model(size, problems, rng):
    """
    对一个规模的模型，在所有题目上跑所有做法。

    【返回】(results, stats)：results 是 {做法: 准确率}；stats 记录 CoT / 直接回答分支的正确率和平均 Δ。
    """
    model_cot, model_direct = ToyLLM(size, cot=True), ToyLLM(size, cot=False)
    p_first = first_token_probs()
    hits = {name: [] for name in METHODS}
    stats = {"cot": {"correct": [], "delta": []}, "direct": {"correct": [], "delta": []}}
    for problem in problems:
        answer = problem.answer()
        kinds = branch_types(rng)
        probe = direct_distribution(model_direct, problem, rng)
        paths = [decode_branch(k, model_cot, model_direct, problem, probe, rng) for k in kinds]
        for k, (a, ok, d) in zip(kinds, paths):
            stats[k]["correct"].append(ok)
            stats[k]["delta"].append(d)

        hits["贪心（不提示）"].append(paths[0][1])
        longest = kinds.index("cot") if "cot" in kinds else 0          # 有推理的路径最长；没有就退回贪心
        hits["按长度挑（第一条有推理的路径）"].append(paths[longest][1])
        best = max(range(K_BRANCH), key=lambda i: paths[i][2])
        hits["CoT-decoding（最大 Δ）"].append(paths[best][1])
        totals = {}
        for a, _, d in paths:
            totals[a] = totals.get(a, 0.0) + d
        hits["CoT-decoding（按答案聚合 Δ）"].append(max(totals, key=lambda a: totals[a]) == answer)
        sc = [decode_branch(kinds[sample_index(p_first, rng)], model_cot, model_direct, problem, probe, rng)[0]
              for _ in range(N_SC)]
        hits["自洽性（不提示，10 个样本）"].append(majority_vote(sc)[0] == answer)
        hits["CoT 提示（1 个样本）"].append(model_cot.generate(problem, rng, temperature=TEMPERATURE).correct)
        votes = [model_cot.generate(problem, rng, temperature=TEMPERATURE).answer for _ in range(N_SC)]
        hits["CoT 提示 + 自洽性（10 个样本）"].append(majority_vote(votes)[0] == answer)
        hits["K 个分支的覆盖率（oracle）"].append(any(ok for _, ok, _ in paths))
    return {name: mean(hits[name]) for name in METHODS}, stats


METHODS = ("贪心（不提示）", "按长度挑（第一条有推理的路径）", "CoT-decoding（最大 Δ）", "CoT-decoding（按答案聚合 Δ）", "自洽性（不提示，10 个样本）",
           "CoT 提示（1 个样本）", "CoT 提示 + 自洽性（10 个样本）", "K 个分支的覆盖率（oracle）")


def main():
    problems = make_problem_set(random.Random(120), N_PROBLEMS, 3, 6)
    sizes = (1e9, 1e10, 1e11)
    table, all_stats = {}, {}
    for size in sizes:
        table[size], all_stats[size] = run_model(size, problems, random.Random(int(1200 + math.log10(size))))

    print("=" * 78)
    print(f"准确率（{N_PROBLEMS} 道 3~6 步的题；分支之后 temperature = {TEMPERATURE}）")
    print("=" * 78)
    print(pad("", 34) + "".join(pad(fmt_size(s), 10, "right") for s in sizes))
    for name in METHODS:
        print(pad(name, 34) + "".join(f"{table[s][name]:10.3f}" for s in sizes))

    print("\n分支的诊断（每种规模的全部分支）：")
    print(pad("", 30) + "".join(pad(fmt_size(s), 10, "right") for s in sizes))
    for kind, label in (("cot", "CoT 分支"), ("direct", "直接回答分支")):
        print(pad(label + "：正确率", 30) + "".join(f"{mean(all_stats[s][kind]['correct']):10.3f}" for s in sizes))
        print(pad(label + "：平均 Δ", 30) + "".join(f"{mean(all_stats[s][kind]['delta']):10.3f}" for s in sizes))
    for ok, label in ((True, "答对的 CoT 分支：平均 Δ"), (False, "答错的 CoT 分支：平均 Δ")):
        row = ""
        for s in sizes:
            st = all_stats[s]["cot"]
            ds = [d for d, c in zip(st["delta"], st["correct"]) if c == ok]
            row += f"{mean(ds):10.3f}" if ds else pad("—", 10, "right")
        print(pad(label, 30) + row)

    print("\n观察：")
    print("  · 贪心解码走的是直接回答，在多步的题上很差（1B 只有 0.055）。只要在第一个 token 处多看几个候选，")
    print("    按答案的置信度 Δ 挑一条（CoT-decoding），准确率就大幅提高（1B：0.055 → 0.265）——")
    print("    不需要任何提示，推理路径本来就在模型的解码空间里。")
    print("  · Δ 有用，因为它和\"有没有推理\"、\"推理对不对\"都相关：CoT 分支的平均 Δ 远高于直接回答分支，")
    print("    答对的 CoT 分支又高于答错的。但 Δ 只看答案 token，看不到前面步骤里的错误，所以只比")
    print("    \"随便挑一条有推理的路径\"（按长度挑）好一点点。")
    print("  · 把答案相同的路径的 Δ 加起来（聚合）明显更好（10B：0.540 → 0.675）：这就是 self-consistency 的")
    print("    \"对推理路径求边缘分布\"——我们要的是最可能的**答案**，而不是最可能的**推理 + 答案**。")
    print("  · 同样是 10 条路径，聚合后的 CoT-decoding 好于不提示的自洽性；在 1B、10B 上仍不如\"CoT 提示 + 自洽性\"")
    print("    （提示让 10 条路径都先推理），100B 时两者打平。分支本身的覆盖率（oracle）高得多，挑选仍是瓶颈。\n")


if __name__ == "__main__":
    main()