"""
=============================================================================
第 8 讲 · 02：Search-o1 —— 推理到一半，不确定了就去搜索；搜到的文档先提炼再用
=============================================================================

对应讲义：lectures/lecture08_search_deep_research.md 第 4 节（Li et al. 2025）

-----------------------------------------------------------------------------
【设定】
-----------------------------------------------------------------------------
复用第 4 讲 01 号文件的小世界：题目里有两个数字是**事实**（"仓库 07 的库存"），要先知道它们的值，
再做 3~4 步算术。事实有流行度；模型的参数记忆对冷门事实经常记错，而且总是编出同一个错误的值。
推理模型用 100B 的 ToyLLM（相当于一个 reasoning model：算术比较可靠，知识缺口成为主要问题）。

新增的两个要素：
    · 模型知道自己对一个事实有多大把握（confidence）。这里假设它是**校准**的：把握 = 它真正记对的概率。
    · 搜索返回的是一篇**文档**：除了要找的事实，还提到 3 个相似实体的数字（干扰信息）。
      直接把原始文档塞进推理链，模型有 25% 的概率取错数字（被干扰）；
      Reason-in-Documents 先读文档、结合当前的问题提炼出需要的那一句，再放进推理链，取错的概率只有 3%。

五种做法：
    直接推理         ：不检索，事实全靠记忆。
    标准 RAG         ：推理之前，用整个问题检索一次，拿回 2 篇文档；每个需要的事实有 75% 的概率被覆盖；
                      原始文档直接放进上下文；没覆盖到的事实靠记忆。
    agentic RAG     ：推理中途，对把握不足（< 阈值）的事实主动搜索；原始文档直接放进推理链。
    Search-o1       ：同上，但搜到的文档先经过 Reason-in-Documents 提炼。
    Search-o1（总是搜）：每个事实都搜（阈值 = 1）。

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

import importlib.util
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"))

# 【Python 小知识】第 4 讲的文件名以数字开头（01_react.py），不能直接 import。
#   importlib.util.spec_from_file_location 可以按文件路径加载一个模块，然后像普通模块一样使用它的函数。
_spec = importlib.util.spec_from_file_location("react_world", os.path.join(_HERE, "..", "lecture04", "01_react.py"))
react_world = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(react_world)

from toy_llm import ToyLLM, mean, pad  # noqa: E402

DISTRACT_RAW, DISTRACT_REFINED = 0.25, 0.03
RAG_COVERAGE = 0.75
THRESHOLD = 0.9


def confidence(model, fact):
    """
    模型对一个事实的把握：和第 4 讲 recall 里"记对的概率"是同一个公式（假设模型是校准的）。

    【边界】真实的 LLM 往往**不是**校准的：对记错的冷门事实也很自信（第 4 讲 CoT-SC → ReAct 的教训）。
        这里假设校准，是为了看清"如果模型知道自己不知道，按需搜索能带来什么"。
    """
    z = -2.5 + math.log10(fact["popularity"]) + 0.8 * math.log10(model.size / 1e9)
    return 1 / (1 + math.exp(-z))


def read_document(facts, fact_id, rng, refined):
    """
    搜索一个事实并从文档里读出它的值。

    【思路】先用第 4 讲的 search 工具查（85% 查到、12% 没结果、3% 查到相似实体）；
        查到了，再从文档里"读"出数字：原始文档有 DISTRACT_RAW 的概率读成干扰项（另一个实体的数字），
        Reason-in-Documents 提炼后只有 DISTRACT_REFINED 的概率读错。
    【返回】读出的值，或 None（没有结果）
    """
    v = react_world.search(facts, fact_id, rng)
    if v is None:
        return None
    if rng.random() < (DISTRACT_REFINED if refined else DISTRACT_RAW):
        return facts[(fact_id + 1 + rng.randrange(3)) % len(facts)]["value"]   # 读成了干扰项
    return v


def solve(method, model, fp, facts, rng, threshold=THRESHOLD):
    """
    用某种做法解一道题，返回 (答案, 搜索次数)。

    【参数】method："direct" / "rag" / "agentic" / "search_o1"；threshold：把握低于它就去搜
    """
    values, n_search = [], 0
    for i in fp["fact_ids"]:
        fact = facts[i]
        v = None
        if method == "rag":
            if rng.random() < RAG_COVERAGE:                     # 这个事实被事先检索的文档覆盖到了
                v = read_document(facts, i, rng, refined=False)
        elif method in ("agentic", "search_o1") and confidence(model, fact) < threshold:
            n_search += 1
            v = read_document(facts, i, rng, refined=(method == "search_o1"))
        if v is None:
            v = react_world.recall(model, fact, i, rng)
        values.append(v)
    if method == "rag":
        n_search = 1
    return model.generate(react_world.build_problem(fp, values), rng, 0.7).answer, n_search


def main():
    rng = random.Random(200)
    facts = react_world.make_facts(random.Random(91))
    problems = [react_world.make_fact_problem(rng, facts, pid) for pid in range(300)]
    model = ToyLLM(1e11)

    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）")

    print("=" * 78)
    print(f"演示 1：五种做法（{model.name} 推理模型，300 道要查两个事实的题，每题跑 3 次；按需搜索的阈值 = {THRESHOLD}）")
    print("=" * 78)
    rows = (("直接推理", "direct", THRESHOLD), ("标准 RAG", "rag", THRESHOLD), ("agentic RAG", "agentic", THRESHOLD),
            ("Search-o1", "search_o1", THRESHOLD), ("Search-o1（总是搜）", "search_o1", 1.01))
    print(pad("做法", 22) + "".join(pad(t, 18, "right") for t in tiers) + pad("全部", 9, "right")
          + pad("平均搜索次数", 14, "right"))
    for name, method, th in rows:
        by_tier, searches = {t: [] for t in tiers}, []
        for fp in problems:
            truth = react_world.true_answer(fp, facts)
            for _ in range(3):
                ans, n = solve(method, model, fp, facts, rng, th)
                by_tier[tier(fp)].append(ans == truth)
                searches.append(n)
        print(pad(name, 22) + "".join(pad(f"{mean(by_tier[t]):.3f}", 18, "right") for t in tiers)
              + pad(f"{mean(x for t in tiers for x in by_tier[t]):.3f}", 9, "right")
              + pad(f"{mean(searches):.2f}", 14, "right"))

    print("\n" + "=" * 78)
    print("演示 2：什么时候该搜？—— Search-o1 的阈值与准确率、搜索次数")
    print("=" * 78)
    print(pad("阈值（把握低于它就搜）", 24) + pad("准确率", 10, "right") + pad("平均搜索次数", 14, "right"))
    for th in (0.0, 0.5, 0.7, 0.9, 0.97, 1.01):
        acc, searches = [], []
        for fp in problems:
            truth = react_world.true_answer(fp, facts)
            for _ in range(3):
                ans, n = solve("search_o1", model, fp, facts, rng, th)
                acc.append(ans == truth)
                searches.append(n)
        label = "1（总是搜）" if th > 1 else f"{th}"
        print(pad(label, 24) + pad(f"{mean(acc):.3f}", 10, "right") + pad(f"{mean(searches):.2f}", 14, "right"))

    print("\n观察：")
    print("  · 直接推理在长尾事实上很差（知识缺口）。")
    print("  · 标准 RAG 反而比直接推理差，头部的题掉得最多：模型本来就记得的事实，被原始文档里的干扰数字带偏了。")
    print("    检索不是免费的 —— 无关、嘈杂的上下文会损害推理（Shi et al. 2023：LLM 很容易被无关的上下文干扰）。")
    print("    而且只在推理之前检索一次，有些推理中途才发现需要的事实没被覆盖。")
    print("  · agentic RAG 在推理中途按需搜索，长尾提高了；但原始文档同样会干扰，整体仍不如直接推理。")
    print("  · Search-o1 加上 Reason-in-Documents：先结合当前的问题从文档里提炼出需要的那一句，再放进推理链。")
    print("    干扰几乎消失，三个档次都大幅提高。这正是论文设计这个模块的理由：检索到的文档冗长、嘈杂，")
    print("    直接插入会打断推理的连贯性。")
    print("  · 阈值的取舍：在这里搜得越多越准（文档提炼后比记忆可靠），代价是搜索次数。")
    print("    只搜\"把握最低\"的事实，就能用少得多的搜索拿到大部分收益 —— 前提是模型**知道自己不知道**。")
    print("    第 4 讲已经看到，真实的模型对编出来的事实常常很自信，这个前提并不总成立。\n")


if __name__ == "__main__":
    main()