"""
=============================================================================
第 12 讲 · 02：推理的三个局限 —— 前提的顺序、无关的信息、没有外部反馈的自我纠正
=============================================================================

对应讲义：lectures/lecture12_llm_reasoning.md 第 6 节
    · Chen et al. 2024, Premise Order Matters in Reasoning with Large Language Models
    · Shi et al. 2023, Large Language Models Can Be Easily Distracted by Irrelevant Context
    · Huang et al. 2024, Large Language Models Cannot Self-Correct Reasoning Yet

所有演示都用 10B 的 CoT 模型（temperature = 0.7），题目是第 1 讲的多步算术题。
每个演示都在 toy world 里加了一个**设定的机制**（见各演示开头的说明），演示的是机制带来的后果，不是真实模型的数值。

-----------------------------------------------------------------------------
【演示 1：前提的顺序】
-----------------------------------------------------------------------------
题目的各步被**打乱顺序**写出来（每一步标着它是第几步："第 3 步：× 4"）。逻辑上题目完全没变，
但模型要按第 1、2、3……步的顺序计算，每一步都要在题目里"找"到下一个前提。
设定的机制：如果下一个需要的前提不紧跟在上一个后面，这一步更难 ——
    correct 的 logit 减去 GAMMA × min(|位置差 − 1|, 3)（跳得越远越难找，最多罚 3 份）。

-----------------------------------------------------------------------------
【演示 2：无关的信息】
-----------------------------------------------------------------------------
每道题插入 2 句无关的话（"小明的妹妹另外拿来了 23 个，但这和问题无关"）。每句话有一个"迷惑性" c（0~1，
每道题固定）。设定的机制：模型每次生成时，以 sigmoid(6 × (c − 0.85) − 指令的效果) 的概率把这句话
当成题目的一步算进去（迷惑性 c > 0.85 的句子，被算进去的概率超过一半）。提示里加上"忽略无关的信息"，指令的效果 = 1.5；否则为 0。

-----------------------------------------------------------------------------
【演示 3：没有外部反馈的自我纠正】
-----------------------------------------------------------------------------
模型先回答一次，然后被要求"检查你的答案，找出问题并改进"：
    · 内在自我纠正：模型用同样的温度重新算一遍；结果不同就以 80% 的概率改成新结果。
      另外，提示暗示"有问题"，即使两次结果相同，模型也有 PHANTOM = 15% 的概率"找到一个问题"，
      改动其中一步（手滑）—— 设定的机制，对应论文观察到的"把对的改成错的"。
    · 复查比生成更可靠：用更低的温度（0.3）重新算，不被暗示"有问题"。这相当于存在 generation–verification gap
      （第 1、3 讲）。
    · 用标准答案决定要不要改（oracle）：只有答错时才让模型自我纠正。很多早期论文的实验无意中这样做了；
      它用到了部署时拿不到的信息。

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

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, PROCEDURES, Problem, Step, ToyLLM, apply_bug, apply_op, majority_vote,  # noqa: E402
                     make_problem_set, mean, pad, sample_index, softmax)

TEMPERATURE = 0.7
CHECK_TEMPERATURE = 0.3
MODEL = ToyLLM(1e10)
GAMMA = 0.2                 # 演示 1：每"跳"一个位置，correct 的 logit 减多少
DISTRACT_SLOPE, DISTRACT_CENTER, INSTRUCTION_EFFECT = 6.0, 0.85, 1.5
SWITCH_P, PHANTOM = 0.8, 0.15


def generate_with_order(model, problem, order, rng, temperature=TEMPERATURE):
    """
    按题目里**写出来的顺序** order 找前提，按逻辑顺序（第 1、2、3……步）计算。

    【参数】order：list，order[j] = 第 j 个写出来的是逻辑上的第几步。原顺序就是 [0, 1, …, k−1]。
    【返回】(答案, 是否正确)
    【思路】和 ToyLLM.generate 一样一步一步地算，只是每一步的 correct logit 多减一项"找前提的难度"：
        上一步的前提写在位置 p_prev，这一步的写在位置 p_now；紧挨着（p_now = p_prev + 1）不罚，
        否则罚 GAMMA × min(|p_now − p_prev − 1|, 3)。第 1 步的"上一个位置"当作 −1（从题目开头读起）。
    【边界】order 必须是 0~k−1 的一个排列，否则 position 字典会缺键（KeyError）。
    """
    position = {logical: written for written, logical in enumerate(order)}
    x, prev = problem.start, -1
    for i, step in enumerate(problem.steps):
        logits = model.step_logits(step.op, x, step.operand, i)
        logits[0] -= GAMMA * min(abs(position[i] - prev - 1), 3)
        proc = PROCEDURES[sample_index(softmax(logits, temperature), rng)]
        right = apply_op(step.op, x, step.operand)
        if proc == "correct":
            x = right
        elif proc == "bug":
            x = apply_bug(step.op, x, step.operand)
        else:
            x = right + CARELESS_OFFSETS[rng.randrange(len(CARELESS_OFFSETS))]
        prev = position[i]
    return x, x == problem.answer()


def kendall_distance(order):
    """
    一个排列离"原顺序"有多远：逆序对的个数（Kendall tau 距离）。0 = 原顺序；k(k−1)/2 = 完全倒过来。

    【Python 小知识】两层循环 for i … for j in range(i + 1, n) 枚举所有 i < j 的位置对。
    """
    n = len(order)
    return sum(1 for i in range(n) for j in range(i + 1, n) if order[i] > order[j])


def demo_premise_order():
    print("=" * 78)
    print("演示 1：前提的顺序（6 步的题，200 道，每道题每种顺序采样 10 次）")
    print("=" * 78)
    rng = random.Random(1210)
    problems = make_problem_set(rng, 200, 6, 6)
    groups = {"原顺序": [], "轻微打乱（逆序对 1~4）": [], "大幅打乱（逆序对 5~10）": [], "完全倒序": []}
    for problem in problems:
        k = problem.n_steps
        orders = [list(range(k)), list(range(k - 1, -1, -1))]
        for _ in range(4):
            o = list(range(k))
            rng.shuffle(o)
            orders.append(o)
        for o in orders:
            d = kendall_distance(o)
            if d == 0:
                key = "原顺序"
            elif d == k * (k - 1) // 2:
                key = "完全倒序"
            elif d <= 4:
                key = "轻微打乱（逆序对 1~4）"
            elif d <= 10:
                key = "大幅打乱（逆序对 5~10）"
            else:
                continue                              # 逆序对 11~14 的随机排列不单独统计
            groups[key].extend(generate_with_order(MODEL, problem, o, rng)[1] for _ in range(10))
    for key, hits in groups.items():
        print(pad(key, 28) + f"准确率 {mean(hits):.3f}（{len(hits) // 10} 个题目 × 顺序）")
    print("\n观察：")
    print("  · 题目的逻辑完全相同，只是前提的顺序变了，准确率就明显下降，倒序最差。")
    print("    LLM 是按顺序读、按顺序生成的：前提的顺序和推理的顺序一致时，每一步需要的信息\"就在手边\"；")
    print("    顺序被打乱，每一步都要在上下文里来回找。论文（R-GSM 等）报告，打乱前提顺序可以让准确率下降超过 30%。\n")


def with_distractors(problem, distractors, included):
    """
    把被模型"当真"的无关句子作为额外的步骤插进题目里。

    【参数】distractors：[(插入位置, 运算数), ...]；included：和 distractors 等长的 bool 列表
    【返回】新的 Problem（标准答案仍以原题为准，由调用者比较）
    【Python 小知识】Problem 是 frozen dataclass（不能修改），所以要新建一个。
    """
    steps = list(problem.steps)
    for (pos, operand), inc in sorted(zip(distractors, included), key=lambda t: -t[0][0]):
        if inc:
            steps.insert(pos, Step("add", operand))       # 从后往前插，前面的插入位置不会被影响
    return Problem(problem.pid, problem.start, tuple(steps))


def sigmoid(z):
    return 1.0 / (1.0 + math.exp(-z))


def demo_distraction():
    print("=" * 78)
    print("演示 2：无关的信息（3~6 步的题，200 道，每道题插入 2 句无关的话）")
    print("=" * 78)
    rng = random.Random(1220)
    problems = make_problem_set(rng, 200, 3, 6)
    setups = []
    for problem in problems:
        ds = [(rng.randint(0, problem.n_steps), rng.randint(10, 99)) for _ in range(2)]
        cs = [rng.random() for _ in range(2)]
        setups.append((problem, ds, cs))

    def one_answer(problem, ds, cs, instruction):
        effect = INSTRUCTION_EFFECT if instruction else 0.0
        included = [rng.random() < sigmoid(DISTRACT_SLOPE * (c - DISTRACT_CENTER) - effect) for c in cs]
        return MODEL.generate(with_distractors(problem, ds, included), rng, TEMPERATURE).answer

    rows = []
    rows.append(("没有无关的信息（1 个样本）", mean(MODEL.generate(p, rng, TEMPERATURE).correct for p, _, _ in setups)))
    for instruction, label in ((False, ""), (True, " + 提示\"忽略无关信息\"")):
        single = mean(one_answer(p, ds, cs, instruction) == p.answer() for p, ds, cs in setups)
        sc = mean(majority_vote([one_answer(p, ds, cs, instruction) for _ in range(10)])[0] == p.answer()
                  for p, ds, cs in setups)
        rows.append(("有无关的信息" + label + "（1 个样本）", single))
        rows.append(("有无关的信息" + label + "（自洽性 10 个）", sc))
    baseline_sc = mean(majority_vote([MODEL.generate(p, rng, TEMPERATURE).answer for _ in range(10)])[0] == p.answer()
                       for p, _, _ in setups)
    rows.insert(1, ("没有无关的信息（自洽性 10 个）", baseline_sc))
    for name, acc in rows:
        print(pad(name, 50) + f"{acc:.3f}")

    print("\n按迷惑性拆开（有无关的信息，不加提示）：")
    for has_high, label in ((False, "两句话的迷惑性都 ≤ 0.85"), (True, "至少一句的迷惑性 > 0.85")):
        group = [(p, ds, cs) for p, ds, cs in setups if (max(cs) > DISTRACT_CENTER) == has_high]
        single = mean(one_answer(p, ds, cs, False) == p.answer() for p, ds, cs in group)
        sc = mean(majority_vote([one_answer(p, ds, cs, False) for _ in range(10)])[0] == p.answer()
                  for p, ds, cs in group)
        print(pad(f"{label}（{len(group)} 道）", 40) + f"1 个样本 {single:.3f}    自洽性 10 个 {sc:.3f}")
    print("\n观察：")
    print("  · 加入两句无关的话，准确率大约减半。提示\"忽略无关的信息\"能挽回一大部分，但不是全部。")
    print("  · 自洽性对两种错误的作用不同：迷惑性不高的句子只是有时被算进去（随机的错误），投票能消掉大部分；")
    print("    迷惑性高的句子**大多数时候**都被算进去（系统性的错误），投票几乎帮不上忙 —— 多数票就是错的答案。")
    print("    这又是第 1 讲\"随机错误 vs 系统错误\"的区别：多采样只能对付前者。")
    print("    （拆开统计时重新采样了，所以两组的加权平均和上面的总数略有出入。）\n")


def perturb(problem, rng):
    """
    "找到一个并不存在的问题"：随机挑一步，在这一步的正确结果上手滑一下，后面的步骤照常算对。

    【返回】改过的答案（int）。几乎一定是错的。
    """
    x = problem.start
    j = rng.randrange(problem.n_steps)
    for i, step in enumerate(problem.steps):
        x = apply_op(step.op, x, step.operand)
        if i == j:
            x += CARELESS_OFFSETS[rng.randrange(len(CARELESS_OFFSETS))]
    return x


def self_correct(problem, answer, mode, rng):
    """
    一轮自我纠正。

    【参数】answer：当前的答案；mode："intrinsic" / "low_temp" / "oracle"
    【返回】新的答案
    【思路】
        intrinsic：同样的温度重算；不同就以 SWITCH_P 改；相同也有 PHANTOM 的概率"找到问题"而改错。
        low_temp ：用 CHECK_TEMPERATURE 重算（更可靠）；不同就以 SWITCH_P 改；不会凭空"找到问题"。
        oracle   ：先偷看标准答案 —— 答对了就不动；答错了才按 intrinsic 的方式纠正。
    """
    if mode == "oracle":
        if answer == problem.answer():
            return answer
        mode = "intrinsic"
    temperature = TEMPERATURE if mode == "intrinsic" else CHECK_TEMPERATURE
    new = MODEL.generate(problem, rng, temperature).answer
    if new != answer:
        return new if rng.random() < SWITCH_P else answer
    if mode == "intrinsic" and rng.random() < PHANTOM:
        return perturb(problem, rng)
    return answer


def demo_self_correction():
    print("=" * 78)
    print("演示 3：没有外部反馈的自我纠正（3~6 步的题，2000 道）")
    print("=" * 78)
    rng = random.Random(1230)
    problems = make_problem_set(rng, 2000, 3, 6)
    first = [MODEL.generate(p, rng, TEMPERATURE).answer for p in problems]
    acc0 = mean(a == p.answer() for a, p in zip(first, problems))
    print(pad("", 34) + pad("第 0 轮", 10, "right") + pad("第 1 轮", 10, "right") + pad("第 2 轮", 10, "right")
          + pad("对→错", 9, "right") + pad("错→对", 9, "right"))
    for mode, label in (("intrinsic", "内在自我纠正"), ("low_temp", "复查比生成更可靠（低温重算）"),
                        ("oracle", "用标准答案决定要不要改（oracle）")):
        answers, accs, to_wrong, to_right = list(first), [acc0], 0, 0
        for _ in range(2):
            new = [self_correct(p, a, mode, rng) for p, a in zip(problems, answers)]
            for p, a, b in zip(problems, answers, new):
                to_wrong += a == p.answer() and b != p.answer()
                to_right += a != p.answer() and b == p.answer()
            answers = new
            accs.append(mean(a == p.answer() for a, p in zip(answers, problems)))
        print(pad(label, 34) + "".join(f"{x:10.3f}" for x in accs) + f"{to_wrong:9d}{to_right:9d}")
    print("\n观察：")
    print("  · 内在自我纠正没有带来提升，反而一轮比一轮差：模型用和生成时一样不可靠的方式\"复查\"，")
    print("    把错的改对的同时，也把对的改错了；提示暗示\"有问题\"，让它更容易改错。")
    print("    Huang et al.：GPT-3.5 在 GSM8K 上自我纠正两轮，75.9% → 74.7%；在 CommonSenseQA 上 75.8% → 38.1%（第 1 轮）。")
    print("  · 复查本身更可靠时（低温重算，相当于存在 generation–verification gap），自我纠正才有用 ——")
    print("    第 2 讲的 revise 能提升，正是因为这个 gap，而不是因为\"自我纠正\"本身。")
    print("  · 用标准答案决定\"要不要改\"，准确率单调上升（对的永远不会被改错）。这个提升是假的：")
    print("    部署时没有标准答案。早期一些自我纠正的论文无意中这样做了，结论因此过于乐观。\n")


if __name__ == "__main__":
    demo_premise_order()
    demo_distraction()
    demo_self_correction()