"""
=============================================================================
第 4 讲 · 01：ReAct —— 交替地思考和行动，用工具的返回结果纠正推理
=============================================================================

对应讲义：lectures/lecture04_feedback_tools_code.md 第 2 节（Yao et al. 2022）

-----------------------------------------------------------------------------
【这个文件的 toy world：要先查资料的算术题】
-----------------------------------------------------------------------------
题目里的一些数字不是直接给出的，而是**事实**，要先查出来：
        从【仓库 07 的库存】开始：+【工厂 23 的日产量】，× 6，- 48。最后得多少？
每个事实有一个"流行度"（在预训练数据里出现的次数，1 ~ 10000）。模型可以用两种方式得到事实的值：

    · 参数记忆（memory）：直接"回忆"。流行的事实记得准；冷门的事实（长尾）经常记错，
      而且错得很"固定"——模型会一次又一次地编出同一个错误的值（confabulation），并且很自信。
    · 搜索工具（search）：85% 返回正确的值，12% 查不到（"没有结果"），3% 查到一个相似但错误的事实。

拿到所有数字后，再用第 1 讲的 ToyLLM 做算术（CoT 或 direct）。

-----------------------------------------------------------------------------
【五种 agent】
-----------------------------------------------------------------------------
    CoT         只推理、不行动：事实全靠记忆，算术用 CoT。
    CoT-SC      CoT 采样 5 次，对答案做 majority vote（self-consistency）。
    Act         只行动、不推理：用搜索查事实，但没有"想一想"的过程——查不到就放弃（直接猜一个数），
                算术也不写中间步骤（direct 模式）。
    ReAct       交替思考和行动：Thought（我需要查仓库 07 的库存）→ Action（search）→ Observation（结果）→
                Thought（没查到，换个说法再查）→ ……；算术用 CoT。每次思考有小概率"陷入循环"（重复同样的动作），
                用完步数上限就交不出答案。
    组合        ReAct → CoT-SC：ReAct 交不出答案时，退回 CoT-SC；
                CoT-SC → ReAct：CoT-SC 的 5 个答案里最常见的那个不到一半（没把握）时，改用 ReAct。
                这是 ReAct 论文在 HotpotQA、FEVER 上效果最好的做法。

【重要】这些 agent 的行为方式（谁会重试、谁写 CoT、循环的概率）是按 ReAct 论文描述的现象**设定**的，
所以这个文件演示的是"这些机制怎样组合出论文里的结论"，而不是独立地"发现"这些结论。

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

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 Problem, Step, ToyLLM, majority_vote, mean, pad  # noqa: E402

N_FACTS = 60
SEARCH_OK, SEARCH_EMPTY = 0.85, 0.12          # 剩下的 3%：查到一个相似但错误的事实
REACT_MAX_RETRIES = 2                          # ReAct 查不到时最多换说法重查 2 次
REACT_LOOP_PROB = 0.04                         # ReAct 每次"思考"陷入循环的概率
CONFAB_FIXED = 0.8                             # 记错时，80% 编出"固定的"错误值，20% 随机编一个


# =============================================================================
# 事实库与题目
# =============================================================================
def make_facts(rng):
    """
    生成事实库：每个事实有名字、真实的值（10~99）和流行度（1~10000，对数均匀）。

    【返回】list of dict：{"name", "value", "popularity"}
    【思路】对数均匀：流行度 = 10^U，U 在 [0, 4] 上均匀。这样"很冷门"和"很流行"的事实一样多 ——
        真实世界的知识分布是长尾的（CS224V 第 1 讲），大量事实只出现过几次。
    """
    kinds = ("仓库", "工厂", "车队", "门店", "机房")
    facts = []
    for i in range(N_FACTS):
        name = f"{kinds[i % len(kinds)]} {i:02d} 的{('库存', '日产量', '车辆数', '客流', '服务器数')[i % 5]}"
        facts.append({"name": name, "value": rng.randint(10, 99), "popularity": 10 ** rng.uniform(0, 4)})
    return facts


def make_fact_problem(rng, facts, pid):
    """
    生成一道"要先查资料"的题：起始数和一个加减法的运算数是事实，其余是普通的运算。

    【返回】dict：{"pid", "fact_ids"（两个事实的下标）, "steps"（Step 列表，第 0 步的运算数要换成事实的值）}
    """
    a, b = rng.sample(range(len(facts)), 2)
    steps = [Step(rng.choice(("add", "sub")), -1)]               # 运算数 −1 表示"这里填事实 b 的值"
    for _ in range(rng.randint(2, 3)):
        op = rng.choice(("add", "sub", "mul"))
        steps.append(Step(op, rng.randint(2, 15) if op == "mul" else rng.randint(10, 99)))
    return {"pid": pid, "fact_ids": (a, b), "steps": steps}


def build_problem(fp, values):
    """
    用"agent 认为的"两个事实的值，构造出一道普通的 toy_llm.Problem。

    【思路】如果 agent 查错了或记错了事实，构造出来的就是一道"错误的题"：后面的算术即使全对，
        答案也是错的 —— 这就是幻觉（hallucination）怎样传播到最终答案的。
    """
    first = Step(fp["steps"][0].op, values[1])
    return Problem(fp["pid"], values[0], tuple([first] + list(fp["steps"][1:])))


def true_answer(fp, facts):
    """用真实的事实值算出的标准答案（上帝视角）。"""
    a, b = fp["fact_ids"]
    return build_problem(fp, (facts[a]["value"], facts[b]["value"])).answer()


# =============================================================================
# 记忆与工具
# =============================================================================
def recall(model, fact, fact_id, rng):
    """
    用参数记忆回忆一个事实的值。

    【思路】记对的概率 = σ(−2.5 + log10(流行度) + 0.8 × log10(参数量 / 1e9))：
        越流行、模型越大，越容易记对。记错时，80% 编出一个**固定的**错误值（同一个模型对同一个事实总是编同一个数），
        20% 随机编一个。模型不会说"我不知道"——它总是给出一个数。
    【参数】model：ToyLLM；fact：事实（dict）；fact_id：下标；rng
    【返回】int（回忆出的值，可能是错的）
    """
    z = -2.5 + math.log10(fact["popularity"]) + 0.8 * math.log10(model.size / 1e9)
    if rng.random() < 1 / (1 + math.exp(-z)):
        return fact["value"]
    if rng.random() < CONFAB_FIXED:
        confab = random.Random(f"confab|{fact_id}|{model.size}")
        wrong = confab.randint(10, 98)
        return wrong + 1 if wrong >= fact["value"] else wrong     # 保证编出来的值 ≠ 真实值
    return rng.randint(10, 99)


def search(facts, fact_id, rng):
    """
    搜索工具：85% 返回正确的值；12% 返回 None（没有结果）；3% 返回另一个事实的值（查错了）。
    """
    u = rng.random()
    if u < SEARCH_OK:
        return facts[fact_id]["value"]
    if u < SEARCH_OK + SEARCH_EMPTY:
        return None
    return facts[(fact_id + 1) % len(facts)]["value"]


# =============================================================================
# 五种 agent
# =============================================================================
def agent_cot(model, fp, facts, rng):
    """CoT：事实靠记忆，算术用 CoT。返回答案（int）。"""
    values = [recall(model, facts[i], i, rng) for i in fp["fact_ids"]]
    return model.generate(build_problem(fp, values), rng, 0.7).answer


def agent_cot_sc(model, fp, facts, rng, k=5):
    """CoT-SC：CoT 采样 k 次，返回 (多数答案, 多数答案的得票比例)。"""
    answers = [agent_cot(model, fp, facts, rng) for _ in range(k)]
    answer, votes = majority_vote(answers)
    return answer, votes / k


def agent_act(model, fp, facts, rng):
    """
    Act：只行动、不推理。每个事实搜一次；查不到就放弃，凭记忆猜一个数（没有"想一想、换个说法"的过程）；
    算术不写中间步骤（direct 模式）。
    """
    values = []
    for i in fp["fact_ids"]:
        v = search(facts, i, rng)
        values.append(v if v is not None else recall(model, facts[i], i, rng))
    direct = ToyLLM(model.size, cot=False)
    return direct.generate(build_problem(fp, values), rng, 0.7).answer


def react_lookup(fp, facts, rng):
    """
    ReAct 查资料的部分：Thought → Action → Observation 的循环，返回查到的两个事实的值，或者 None（交不出答案）。

    【思路】对每个事实：
        Thought："我需要查 X" → Action：search(X) → Observation：结果。
        查不到时，Thought："没有结果，换个说法再查" → 最多重查 REACT_MAX_RETRIES 次。
        每次 Thought 都有 REACT_LOOP_PROB 的概率"陷入循环"（重复同样的动作直到步数用完）。
        陷入循环、或者所有重查都失败 → 返回 None。
    【边界】搜索查错时（3%），ReAct 并不知道，会带着错误的值继续往下算。
    """
    values = []
    for i in fp["fact_ids"]:
        got = None
        for _attempt in range(1 + REACT_MAX_RETRIES):
            if rng.random() < REACT_LOOP_PROB:
                return None                               # 陷入循环，步数用完
            got = search(facts, i, rng)
            if got is not None:
                break
        if got is None:
            return None                                   # 换了几种说法都查不到
        values.append(got)
    return tuple(values)


def agent_react(model, fp, facts, rng):
    """
    ReAct：先用 react_lookup 查到事实（Thought / Action / Observation），再 Thought："现在可以算了"，用 CoT 做算术。

    【返回】答案（int）或 None（交不出答案）
    """
    values = react_lookup(fp, facts, rng)
    if values is None:
        return None
    return model.generate(build_problem(fp, values), rng, 0.7).answer


def agent_react_then_cot_sc(model, fp, facts, rng):
    """ReAct → CoT-SC：ReAct 交不出答案时，退回 CoT-SC。"""
    answer = agent_react(model, fp, facts, rng)
    return answer if answer is not None else agent_cot_sc(model, fp, facts, rng)[0]


def agent_cot_sc_then_react(model, fp, facts, rng):
    """CoT-SC → ReAct：CoT-SC 没把握（多数答案得票不到一半）时，改用 ReAct；ReAct 也交不出时用 CoT-SC 的答案。"""
    answer, share = agent_cot_sc(model, fp, facts, rng)
    if share >= 0.5:
        return answer
    alt = agent_react(model, fp, facts, rng)
    return alt if alt is not None else answer


AGENTS = (("CoT", lambda m, fp, f, r: agent_cot(m, fp, f, r)),
          ("CoT-SC", lambda m, fp, f, r: agent_cot_sc(m, fp, f, r)[0]),
          ("Act", agent_act),
          ("ReAct", agent_react),
          ("ReAct→CoT-SC", agent_react_then_cot_sc),
          ("CoT-SC→ReAct", agent_cot_sc_then_react))


def main():
    rng = random.Random(90)
    facts = make_facts(random.Random(91))
    problems = [make_fact_problem(rng, facts, pid) for pid in range(300)]
    # 按两个事实里"较冷门的那个"的流行度，把题目分成三组
    def tier(fp):
        pop = min(facts[i]["popularity"] for i in fp["fact_ids"])
        return "头部（≥ 1000）" if pop >= 1000 else "中间（30~1000）" if pop >= 30 else "长尾（< 30）"
    tiers = ("头部（≥ 1000）", "中间（30~1000）", "长尾（< 30）")
    model = ToyLLM(1e10)

    print("=" * 78)
    print(f"演示 1：五种 agent 的准确率（{model.name} 模型，300 道要查两个事实的 3~4 步题，每题跑 3 次）")
    print("=" * 78)
    counts = {t: sum(1 for fp in problems if tier(fp) == t) for t in tiers}
    print(pad("agent", 15) + "".join(pad(f"{t} {counts[t]}题", 22, "right") for t in tiers)
          + pad("全部", 9, "right") + pad("交不出答案", 12, "right"))
    for name, agent in AGENTS:
        by_tier = {t: [] for t in tiers}
        none_count = 0
        for fp in problems:
            truth = true_answer(fp, facts)
            for _ in range(3):
                ans = agent(model, fp, facts, rng)
                none_count += ans is None
                by_tier[tier(fp)].append(ans == truth)
        overall = mean(x for t in tiers for x in by_tier[t])
        print(pad(name, 15) + "".join(pad(f"{mean(by_tier[t]):.3f}", 22, "right") for t in tiers)
              + pad(f"{overall:.3f}", 9, "right") + pad(f"{none_count / (3 * len(problems)):.1%}", 12, "right"))

    print("\n观察：")
    print("  · CoT 在头部事实上明显好于长尾：长尾事实模型会自信地编出一个错误的值，后面的算术再对也没用。")
    print("    （所有 agent 的绝对准确率都不高，因为 temperature = 0.7 下 10B 模型的多步算术本身就常出错。）")
    print("  · CoT-SC 在头部上提升很大（投票消除了算术上的手滑），在长尾上几乎没用：")
    print("    模型每次都编**同一个**错误的值，投票只会让错误更\"一致\"。")
    print("  · Act 靠搜索在长尾上远好于 CoT，但不会\"想一想\"：查不到就瞎猜，算术不写中间步骤，多步计算经常出错。")
    print("  · ReAct 把两者结合起来：查资料避免了幻觉，写出推理过程改善了算术，长尾上是 CoT 的 4 倍多。")
    print("    但它有自己的失败方式：陷入循环、换了说法也查不到 —— 约 10% 的时候交不出答案。")
    print("  · ReAct → CoT-SC 整体最好：交不出答案时退回 CoT-SC，把这 10% 救回来一部分（提升不大，约 0.02）。")
    print("  · CoT-SC → ReAct 在头部上最好，在长尾上却很差：它靠\"投票有没有把握\"决定是否求助 ReAct，")
    print("    而模型对编出来的值**很有把握**（5 次都编同一个数），所以在最需要查资料的长尾题上反而不去查。")
    print("    \"自己有没有把握\"不是知识缺口的可靠信号 —— 和第 1、3 讲\"自己检查自己\"是同一个问题。")
    print("  · ReAct 论文在 HotpotQA、FEVER 上的结论是：内部知识（CoT）和外部知识（ReAct）互补，两种组合都胜过单一方法。\n")

    # 演示 2：失败原因的分解
    print("=" * 78)
    print("演示 2：CoT 和 ReAct 为什么失败？（长尾和中间组，每题跑 5 次）")
    print("=" * 78)
    hard = [fp for fp in problems if tier(fp) != tiers[0]]
    reasons = {"CoT": {"事实错（幻觉）": 0, "算术错": 0, "正确": 0},
               "ReAct": {"事实错（搜索出错）": 0, "交不出答案": 0, "算术错": 0, "正确": 0}}
    for fp in hard:
        truth = true_answer(fp, facts)
        true_vals = tuple(facts[i]["value"] for i in fp["fact_ids"])
        for _ in range(5):
            vals = tuple(recall(model, facts[i], i, rng) for i in fp["fact_ids"])
            ans = model.generate(build_problem(fp, vals), rng, 0.7).answer
            key = "正确" if ans == truth else ("事实错（幻觉）" if vals != true_vals else "算术错")
            reasons["CoT"][key] += 1
            got = react_lookup(fp, facts, rng)
            if got is None:
                reasons["ReAct"]["交不出答案"] += 1
                continue
            ans = model.generate(build_problem(fp, got), rng, 0.7).answer
            key = "正确" if ans == truth else ("事实错（搜索出错）" if got != true_vals else "算术错")
            reasons["ReAct"][key] += 1
    for name, table in reasons.items():
        total = sum(table.values())
        print(f"  {pad(name, 7)}" + "   ".join(f"{k} {v / total:.1%}" for k, v in table.items()))
    print("\n观察：")
    print("  · CoT 的失败主要来自幻觉：长尾事实记错了，后面的推理建立在错误的前提上。")
    print("    ReAct 论文对 HotpotQA 的错误分析也是如此：CoT 的失败里，幻觉占了大头。")
    print("  · ReAct 几乎没有幻觉，它的失败分散在\"交不出答案\"（循环、查不到）、搜索出错和算术错上。")
    print("    论文发现 ReAct 的主要问题是推理出错（包括重复同样的动作）和搜索返回了没用的信息。")
    print("  · 外部信息的可靠性成了新的瓶颈：工具查错了，agent 很难察觉。这是\"用工具的反馈纠正推理\"的另一面。\n")


if __name__ == "__main__":
    main()