"""
=============================================================================
第 4 讲 · 代码 03：思维链忠实吗？—— 改掉中间步骤，答案跟着变吗
=============================================================================

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

【这个脚本要回答的问题】
    模型写出的思维链，是它**真正**用来得出答案的过程，还是一段"事后的说明"？Lanham et al.（2023）的办法是：
    **动手改**思维链，看答案变不变。在 scratch 格式的模型（"volel . answer driltese"）上做三种干预：
    1. 把中间国家换成一个**错的**国家，让模型从那里接着写：答案跟着错国家走，还是仍然是正确答案？（演示 1）
    2. 不让它写中间步骤，直接从 "answer" 开始（提前作答）；或者把中间步骤换成一个无意义的占位符。（演示 2）
    3. 不做干预时，模型自己写错了中间国家的那些题，它的答案和它写下的（错的）国家一致吗？（演示 3）

【怎么读输出】
    "跟着思维链"：最终答案 = 写下的那个国家的属性；"无视思维链"：最终答案 = 真正那个国家的属性。

【运行】python3 03_faithfulness.py     （读 02 号脚本训练好的缓存；没有的话先训练约 10 秒；约 10 秒）
【依赖】PyTorch；import 同目录的 cot_lab.py，第 1 讲的 console.py。
"""

import os
import random
import sys

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

NO_LEADER = [x for x in COMBOS if x[0] != "leader"]


def wrong_country(c, t, rng):
    """一个错的国家：和 c 不同，而且它的目标属性也和 c 不同（否则看不出答案跟谁走）。"""
    pool = [x for x in WORLD.countries if x.name != c.name and attr(x, t) != attr(c, t)]
    return rng.choice(pool)


def demo1(m, tok):
    """演示 1：把中间国家换成错的。"""
    print_header("演示 1：替模型写下一个错的中间国家，让它接着写（scratch 模型；首都 / 语言做桥的 8 种组合）")
    rows, res = [], {}
    for label, countries in [("SFT 国家", SFT_COUNTRIES), ("测试国家", EVAL_COUNTRIES)]:
        rng = random.Random(0)
        follow = truth = other = n = 0
        ex = None
        for b, t in NO_LEADER:
            for c in countries:
                w = wrong_country(c, t, rng)
                out = ask(m, tok, render_chat(twohop_question(b, t, c)), prefill=f"{w.name} .")
                a = parse_answer(out, "scratch")
                follow += a == attr(w, t)
                truth += a == attr(c, t)
                n += 1
                ex = ex or (twohop_question(b, t, c), out, attr(c, t), attr(w, t))
        other = n - follow - truth
        res[label] = (follow / n, truth / n)
        rows.append([label, pct(follow, n), pct(truth, n), pct(other, n)])
    print_table(["国家", "跟着思维链（答错国家的属性）", "无视思维链（答真正的答案）", "都不是"], rows)
    q, out, g, gw = ex
    print(f"  例子：{q}\n    被改成 → {out!r}（真正的答案 {g}；错国家的属性 {gw}）")
    print_observations([
        f"把中间国家换掉之后，测试国家上 {100 * res['测试国家'][0]:.1f}% 的答案跟着错国家走，只有 {100 * res['测试国家'][1]:.1f}% "
        "仍然是正确答案。这个模型的思维链是**因果上有效**的：它真的在用自己写下的中间结果，而不是先算好答案、再补一段说明。",
        "这正是我们训练它的方式决定的：训练数据里，答案总是由前面写下的国家决定。真实的大模型没有这样的保证——"
        "它们的思维链是从人类写的推理文本和后续训练里学来的，可能是'真正的计算'，也可能是'看起来合理的说明'。",
    ])


def demo2(m, tok):
    """演示 2：提前作答、用占位符代替中间步骤。"""
    print_header("演示 2：不让它写中间步骤（测试国家；训练过 / 没训练过的组合）")
    rows = []
    res = {}
    variants = [("正常（自己写中间国家）", lambda c: ""),
                ("提前作答：直接从 'answer' 开始", lambda c: "answer"),
                ("中间步骤换成占位符 '<unk> .'", lambda c: "<unk> ."),
                ("中间步骤写对（替它写好正确的国家）", lambda c: f"{c.name} .")]
    for label, combos in [("训练过的组合", [x for x in TRAIN_COMBOS if x[0] != "leader"]),
                          ("没训练过的组合", [x for x in TEST_COMBOS if x[0] != "leader"])]:
        for vname, pre in variants:
            ok = n = 0
            for b, t in combos:
                for c in EVAL_COUNTRIES:
                    out = ask(m, tok, render_chat(twohop_question(b, t, c)), prefill=pre(c))
                    ok += parse_answer(out, "scratch") == attr(c, t)
                    n += 1
            res[(label, vname)] = ok / n
            rows.append([label, vname, pct(ok, n)])
    print_table(["组合", "干预", "最终答对"], rows)
    early = res[("没训练过的组合", variants[1][0])]
    fill = res[("没训练过的组合", variants[2][0])]
    norm = res[("没训练过的组合", variants[0][0])]
    print_observations([
        f"没训练过的组合：正常 {100 * norm:.1f}%，提前作答 {100 * early:.1f}%，占位符 {100 * fill:.1f}%。"
        "去掉中间步骤，答案就垮了——思维链里写的东西是必需的，不是装饰。",
        "Lanham et al.（2023）在大模型上做了同样的干预（提前截断、加入错误、换成同样长度的无意义填充），发现结果因任务而异："
        "在一些任务上，截断思维链几乎不改变答案（说明答案并不依赖它），在另一些任务上则大幅改变。他们还发现，"
        "在一些任务上**越大的模型**，答案越不依赖它写出的推理。",
        "'换成无意义的填充也一样好'是一个重要的对照：如果真是这样，思维链的作用只是'多给了几步计算'，而不是它写下的内容。"
        "在这个迷你模型上，填充不能代替内容。",
    ])


def demo3(m, tok):
    """演示 3：自然发生的错误——答案和写下的中间国家一致吗？"""
    print_header("演示 3：模型自己写错了中间国家时，答案跟着谁走？（全部 12 种组合、测试国家，不做任何干预）")
    agree = disagree_truth = neither = wrong_bridge = right_bridge = right_and_right = no_bridge = wrong_leader = 0
    by_name = {c.name: c for c in WORLD.countries}
    for b, t in COMBOS:
        for c in EVAL_COUNTRIES:
            out = ask(m, tok, render_chat(twohop_question(b, t, c)))
            words = out.split()
            x = by_name.get(words[0]) if words else None
            a = parse_answer(out, "scratch")
            if x is None:
                no_bridge += 1
                continue
            if x.name == c.name:
                right_bridge += 1
                right_and_right += a == attr(c, t)
                continue
            wrong_bridge += 1
            wrong_leader += b == "leader"
            if a == attr(x, t):
                agree += 1
            elif a == attr(c, t):
                disagree_truth += 1
            else:
                neither += 1
    print_table(["情形", "题数"], [["写对了中间国家", right_bridge], ["  其中最终答案也对", right_and_right],
                                 ["写错了中间国家", wrong_bridge], ["  其中答案 = 写下的错国家的属性（忠实地错）", agree],
                                 ["  其中答案 = 真正的答案（不忠实地对）", disagree_truth], ["  其中都不是", neither],
                                 ["  其中用领导人做桥的", wrong_leader], ["第一个词不是国名（没写出中间国家）", no_bridge]])
    print_observations([
        f"写错中间国家的 {wrong_bridge} 道题里，{agree} 道的答案和它写下的错国家一致（'忠实地错'），{disagree_truth} 道仍答对了真正的答案。"
        "大多数错误是'第一步错了，第二步忠实地沿着错的方向走'——这也是 self-consistency 能发挥作用的前提：如果每次采样的第一步"
        "错得不一样，多数投票就能把它们投掉（04 号脚本）。",
        f"写错中间国家的 {wrong_bridge} 道题里，{wrong_leader} 道是领导人做桥的（反转诅咒）：模型不知道领导人对应哪个国家，"
        "只能编一个，然后忠实地按编出来的国家答下去。"
        "**一条忠实的思维链并不保证正确**；它的好处是：错误看得见、查得出（第 10 讲的验证链 CoVe）。",
        "Turpin et al.（2023）指出了相反方向的危险：在真实的大模型里，提示里的偏向（例如少样本示例的答案总是 (A)）会改变模型的答案，"
        "而模型写出的思维链却从不提到这个偏向，只给出一个看似合理的理由——思维链可以系统性地'不忠实'。",
    ])


def main():
    m, tok, _ = get_cot_model(fmt="scratch")
    demo1(m, tok)
    demo2(m, tok)
    demo3(m, tok)


if __name__ == "__main__":
    main()