"""
=============================================================================
第 20 讲 · 共享模块：一个“知道真相”的模拟 LLM 世界（llm_sim.py）
=============================================================================

对应讲义：lectures/lecture20_llm_applications.md 第 3、4、6 节（被 02、03 两个脚本导入）

-----------------------------------------------------------------------------
【这个模块是什么 What this module is】
-----------------------------------------------------------------------------
本讲不调用任何真实的大语言模型 API。我们搭一个小小的“世界”，其中：

    1. **知识库是真实的文本**：一门**虚构**课程「DL-101 深度学习实践」2026 年秋季学期的事务手册
       （作业截止日期、助教答疑时间、迟交政策、成绩构成……）。课程、人名、教室、日期全部是编的。
       `build_world()` 生成 8 节手册、72 个事实，每个事实配一个“直接问法”和一个“改写问法”的问题。
    2. **检索是真实的计算**：中文分词（字的二元组 bigram）、切块（6 种切法）、BM25 和 TF-IDF
       都用 numpy 从零实现。recall@k 是真的算出来的，检索失败就是真的失败。
    3. **“LLM”是一个行为写明的模拟器**：`answer_question()` 按 PARAMS 里写下的规则回答问题——
       它有“参数记忆”（见过去年的手册：有的事实没变、有的已经过时、今年新增的它根本不知道，不知道时
       会编一个看起来合理的答案）；给了资料时会去读资料，但读的可靠性和信息在上下文里的位置有关
       （中间最容易被忽略），资料里同类的事实越多越容易张冠李戴……

**PARAMS 里的每一个概率都是设计者（我）设定的参数，不代表任何真实 LLM 的性能。**
真实模型的这些数字要在你自己的任务上测出来。本讲要学的是**测量和归因的方法**：
怎样算 recall@k、怎样把端到端的错误拆到检索 / 阅读 / 记忆上、怎样读 trace。
正因为模拟器的规则是已知的，我们才能检查这些方法有没有把错误归到正确的地方。

-----------------------------------------------------------------------------
【主要函数】
-----------------------------------------------------------------------------
    build_world(seed)            生成手册（节 → 段落 → 句子）、事实表、问题集
    tokenize(text)               中文按字二元组、英文 / 数字按整词切分
    make_chunks(world, how)      6 种切块方式：按句、固定 2 句、固定 2 句重叠 1 句、按段落、固定 6 句、整节
    BM25Index                    numpy 实现的 BM25 和 TF-IDF 余弦检索
    supports(chunk, fact)        这个块是否包含回答该事实需要的全部句子
    answer_question(...)         模拟 LLM：给定（可能为空的）上下文，回答一个事实问题，并返回一份 trace
    memory_answer(...)           模拟 LLM 的“闭卷”回答（参数记忆）
    hallucinate(...)             按事实的取值类型编一个“看起来合理”的值

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 llm_sim.py        # 自检：打印世界的统计、一段手册、分词和检索的例子（< 1 秒）

只依赖 numpy 和 Python 标准库。
"""

import re

import numpy as np

# =============================================================================
# 0. 模拟器的参数（全部是设计者设定的，不代表任何真实模型）
# =============================================================================
PARAMS = {
    # —— 参数记忆（闭卷）——
    "p_recall": 0.92,            # 记得的事实（没变的 / 过时的），有 92% 的时候能把记忆里的值说出来，否则现编一个
    "p_abstain_unknown": 0.10,   # 完全不知道的事实：10% 的时候老实说“不知道”，90% 的时候编一个看起来合理的答案
    # —— 读资料 ——
    "r_max": 0.97,               # 信息在上下文开头或结尾时，读对的概率
    "mid_drop": 0.30,            # 信息在上下文正中间、且上下文足够长时，读对的概率最多下降多少（"lost in the middle"）
    "l_ref": 30,                 # 上下文有多少句时，中间的下降达到最大（越短的上下文，中间下降越小）
    "conf_per_distractor": 0.012,  # 上下文里每多一个“同类”的其他事实，张冠李戴的概率加多少
    "conf_max": 0.30,            # 张冠李戴概率的上限
    "p_conflict": 0.10,          # 资料和记忆冲突时（记忆里是过时的值），10% 的时候“相信自己的记忆”
    # —— 资料里没有答案时 ——
    "p_grab": {"plain": 0.50, "grounded": 0.15},   # 上下文里有同类的值（但不是要问的那个）：照抄其中一个的概率
    "p_abstain_grounded": 0.85,  # 提示词要求“资料里没有就说没有”时，拒答的概率；否则退回到参数记忆
}

ABSTAIN = "（拒答）"   # 模型说“资料中没有 / 我不知道”时的统一记号

# 记忆状态的中文名（打印表格用）
MEMORY_NAMES = {"known": "没变（记得）", "outdated": "已过时", "unknown": "今年新增"}

# =============================================================================
# 1. 世界：一门虚构课程的事务手册
# =============================================================================
HW_TOPICS = ["逻辑回归", "浅层网络", "深层网络", "正则化", "优化算法", "批归一化",
             "卷积网络", "残差网络", "目标检测", "循环网络", "注意力机制", "Transformer"]
TA_NAMES = ["林晓雨", "周明哲", "陈思远", "王一帆", "赵若溪", "孙浩然", "刘佳宁", "吴子墨"]
ROOMS = ["明德楼 B204", "格物楼 310", "致知楼 105", "博学楼 412", "笃行楼 201", "求是楼 118",
         "慎思楼 307", "明辨楼 220", "弘毅楼 509", "立雪楼 102"]
WEEKDAYS = ["周一", "周二", "周三", "周四", "周五"]
TIME_SLOTS = ["10:00–11:30", "13:30–15:00", "15:00–16:30", "19:00–20:30"]
GRADE_PARTS = ["作业", "测验", "期中考试", "课程项目"]
GRADE_PARTS_ALT = {"作业": "平时作业", "测验": "小测", "期中考试": "期中", "课程项目": "大作业项目"}

# 每种取值类型：怎样把一个数格式化成答案字符串、“编一个”时从哪个范围里抽
VALUE_POOLS = {
    "percent": [f"{x}%" for x in range(5, 45, 5)],
    "days": [f"{x} 天" for x in range(1, 15)],
    "people": [f"{x} 人" for x in range(2, 7)],
    "dollars": [f"{x} 美元" for x in (100, 150, 200, 250, 300, 400, 500)],
    "pages": [f"{x} 页" for x in range(4, 13)],
    "sheets": [f"{x} 页" for x in range(1, 5)],
    "minutes": [f"{x} 分钟" for x in (15, 20, 30, 40, 45, 60, 90)],
    "weekday": list(WEEKDAYS),
    "weekday_time": [f"{d} {t}" for d in WEEKDAYS for t in TIME_SLOTS],
    "room": list(ROOMS),
    "topic": list(HW_TOPICS),
}


def fmt_date(month, day):
    """把 (月, 日) 写成 "10 月 23 日"。【返回】字符串"""
    return f"{month} 月 {day} 日"


def random_date(rng):
    """
    在学期里（9–12 月，每月 1–28 日）随机抽一个日期，用来“编”一个看起来合理的日期。

    【参数】rng：np.random.Generator
    【返回】(month, day)
    【numpy 小知识】rng.integers(a, b) 返回 [a, b) 里的随机整数（不含 b）。
    """
    return int(rng.integers(9, 13)), int(rng.integers(1, 29))


def pool_for(vtype):
    """
    取值类型 vtype 的“候选池”：模型编答案时从这里抽。

    【参数】vtype：取值类型，如 "date"、"percent"
    【返回】字符串列表；"date" 类型返回 None（日期另外用 random_date 生成）
    """
    return VALUE_POOLS.get(vtype)


def hallucinate(vtype, u):
    """
    编一个“看起来合理”的值（幻觉 hallucination）。

    【思路】真实的 LLM 不知道答案时，常常不会说“不知道”，而是给出一个格式正确、数值合理、但没有依据的答案
            ——这正是幻觉危险的地方：它看起来和正确答案一模一样。这里从同类型的候选池里随机抽一个。
    【参数】vtype：取值类型；u：一个 [0, 1) 的随机数（用传进来的随机数而不是自己抽，
            是为了让“同一道题、同一个随机流”在不同配置下用到同样的随机数，见 answer_question 的说明）
    【返回】字符串
    【边界】可能碰巧抽中正确答案（“蒙对了”）——评估时会被算作正确，这和真实评估一样。
    """
    if vtype == "date":
        # 用 u 构造一个确定的日期：把 [0,1) 均匀切成 4 × 28 份
        idx = int(u * 4 * 28)
        return fmt_date(9 + idx // 28, 1 + idx % 28)
    pool = pool_for(vtype)
    return pool[int(u * len(pool))]


def _other_value(rng, vtype, current):
    """从同类型的候选里抽一个和 current 不同的值（生成“去年的旧值”用）。"""
    if vtype == "date":
        # 旧日期 = 新日期前后挪 2–9 天（真实的课程表每年会小幅调整）
        m, d = [int(x) for x in re.findall(r"\d+", current)]
        shift = int(rng.integers(2, 10)) * (1 if rng.random() < 0.5 else -1)
        d2 = min(28, max(1, d + shift))
        if d2 == d:
            d2 = d + 1 if d < 28 else d - 1
        return fmt_date(m, d2)
    pool = [v for v in pool_for(vtype) if v != current]
    return pool[int(rng.integers(len(pool)))]


def build_world(seed=2026):
    """
    生成虚构课程 DL-101 的事务手册、事实表和问题集。

    【思路】手册分 8 节（课程概况、成绩构成、迟交与重新评分、作业安排、测验、助教与答疑、课程项目、考试），
            每节由若干段落组成，每个段落由 1–4 个句子组成。每个“事实”（比如“第 7 次作业的截止日期”）记下：
              · 现在的值 value，以及去年手册里的值 old（模型的“参数记忆”只见过去年的手册）；
              · memory："known"（没变，模型记得）/"outdated"（变了，模型记着旧值）/"unknown"（今年新增，模型没见过）；
              · support：回答它需要的句子编号。注意“作业安排”一节故意写成两句话：
                  “第 7 次作业的主题是卷积网络，……。这次作业的截止日期是 10 月 23 日……。”
                第二句用了代词“这次作业”，单独拿出来就不知道说的是哪次作业——所以回答截止日期需要**两句都在**。
                这是为了让“切块大小”有真实的影响（讲义 §3.5）。
              · 一个直接问法的问题（和手册用词一致）和一个改写问法的问题（换了说法，关键词重合少）。
    【参数】seed：随机种子（决定日期、教室、哪些事实过时等）
    【返回】world 字典：
              "sections"  : 列表，每个元素 {"title", "paragraphs": [[sid, ...], ...]}
              "sentences" : 列表，第 sid 个元素 {"sid", "section", "para", "text"}
              "facts"     : 列表，每个元素是一个事实字典（见上）
              "questions" : 列表，每个元素 {"qid", "fact", "style": "direct"/"para", "text"}
    【边界】同一个 seed 永远生成同一个世界（只用了传进来的 rng，没有用全局随机状态）。
    【Python 小知识】函数里定义的小函数（如下面的 add_sentence）可以读写外层函数的列表变量，
            这叫闭包（closure）；这里用它把“加一句话”的重复代码收起来。
    """
    rng = np.random.default_rng(seed)
    sections, sentences, facts = [], [], []

    def new_section(title, header):
        sections.append({"title": title, "paragraphs": []})
        new_paragraph()
        return add_sentence(header)

    def new_paragraph():
        sections[-1]["paragraphs"].append([])

    def add_sentence(text):
        sid = len(sentences)
        sec = len(sections) - 1
        para = len(sections[-1]["paragraphs"]) - 1
        sentences.append({"sid": sid, "section": sec, "para": para, "text": text})
        sections[-1]["paragraphs"][-1].append(sid)
        return sid

    def add_fact(kind, entity, vtype, value, support, q_direct, q_para, memory=None, old=None):
        """登记一个事实。memory 不给时按 45% / 30% / 25% 随机决定 known / outdated / unknown。"""
        if memory is None:
            r = rng.random()
            memory = "known" if r < 0.45 else ("outdated" if r < 0.75 else "unknown")
        if memory == "known":
            old = value
        elif memory == "outdated" and old is None:
            old = _other_value(rng, vtype, value)
        elif memory == "unknown":
            old = None
        facts.append({"fid": len(facts), "kind": kind, "entity": entity, "vtype": vtype, "value": value,
                      "old": old, "memory": memory, "support": list(support), "value_sid": support[-1],
                      "q_direct": q_direct, "q_para": q_para})

    # ---------------- 第 1 节：课程概况 ----------------
    new_section("课程概况", "【课程概况】DL-101 深度学习实践是一门虚构课程，本手册说明 2026 年秋季学期的安排。")
    new_paragraph()
    lec_day, lec_room = WEEKDAYS[int(rng.integers(5))], ROOMS[int(rng.integers(len(ROOMS)))]
    s = add_sentence(f"课堂讲座每{lec_day}上午 10:30 举行，地点是{lec_room}。")
    add_fact("lec_day", "课堂讲座", "weekday", lec_day, [s], "课堂讲座每周哪天上课？", "这门课的大课是星期几？")
    add_fact("lec_room", "课堂讲座", "room", lec_room, [s], "课堂讲座在哪个教室？", "去哪里听大课？")
    add_sentence("所有课程问题都应该先发在课程论坛上，助教会在两个工作日内回复。")

    # ---------------- 第 2 节：成绩构成 ----------------
    new_section("成绩构成", "【成绩构成】总评由四部分组成，各部分的权重如下。")
    new_paragraph()
    weights = {"作业": 30, "测验": 10, "期中考试": 25, "课程项目": 35}
    for part in GRADE_PARTS:
        s = add_sentence(f"{part}占总评的 {weights[part]}%。")
        add_fact("weight", part, "percent", f"{weights[part]}%", [s], f"{part}占总评的百分之多少？",
                 f"{GRADE_PARTS_ALT[part]}在最终成绩里算多大比重？")

    # ---------------- 第 3 节：迟交与重新评分 ----------------
    new_section("迟交与重新评分", "【迟交与重新评分】请在截止时间前提交，以下规则适用于全部作业。")
    new_paragraph()
    s = add_sentence("每位学生整个学期共有 4 个免费迟交日，可以用在任意作业上。")
    add_fact("late_free", "免费迟交日", "days", "4 天", [s], "每位学生有几个免费迟交日？", "整个学期可以不扣分地晚交几天？")
    s = add_sentence("免费迟交日用完之后，每迟交一天扣该次作业得分的 10%。")
    add_fact("late_pen", "迟交扣分", "percent", "10%", [s], "免费迟交日用完后，每迟交一天扣多少分？",
             "超出宽限以后晚交一天要扣百分之几？")
    s = add_sentence("任何一次作业最多可以迟交 3 天，超过后不再接收。")
    add_fact("late_max", "最多迟交", "days", "3 天", [s], "一次作业最多可以迟交几天？", "作业最晚能拖几天还被接收？")
    new_paragraph()
    s = add_sentence("成绩公布后 7 天内可以在评分系统上申请重新评分。")
    add_fact("regrade", "重新评分", "days", "7 天", [s], "成绩公布后多少天内可以申请重新评分？", "对分数有异议，要在几天之内提出？")

    # ---------------- 第 4 节：作业安排（两句一段，第二句用代词）----------------
    new_section("作业安排", "【作业安排】本学期共有 12 次作业，每次都在太平洋时间 23:59 截止。")
    # 12 个互不相同的截止日期，按作业顺序递增：从 9 月 28 日起的 70 天里抽 12 天
    start = 28                                   # 用“从 9 月 1 日起的第几天”表示日期，9 月 28 日 = 第 28 天
    day_ids = np.sort(rng.choice(np.arange(start, start + 70), size=12, replace=False))
    for k in range(1, 13):
        new_paragraph()
        topic = HW_TOPICS[k - 1]
        s1 = add_sentence(f"第 {k} 次作业的主题是{topic}，附带一份 Jupyter 起始代码。")
        dd = int(day_ids[k - 1])
        month, day = _day_index_to_md(dd)
        due = fmt_date(month, day)
        s2 = add_sentence(f"这次作业的截止日期是 {due}。")
        topic_memory = "known" if rng.random() < 0.8 else "outdated"
        add_fact("hw_topic", f"第 {k} 次作业", "topic", topic, [s1], f"第 {k} 次作业的主题是什么？",
                 f"第 {k} 份作业讲的是哪部分内容？", memory=topic_memory)
        add_fact("hw_due", f"第 {k} 次作业", "date", due, [s1, s2], f"第 {k} 次作业的截止日期是哪天？",
                 f"关于{topic}的那次作业最晚哪天交？")

    # ---------------- 第 5 节：测验（一句一段）----------------
    new_section("测验", "【测验】每次测验都在线上完成，只能提交一次。")
    for k in range(1, 11):
        new_paragraph()
        minutes = VALUE_POOLS["minutes"][int(rng.integers(len(VALUE_POOLS["minutes"])))]
        s = add_sentence(f"第 {k} 次测验在第 {k + 1} 周开放，限时 {minutes}。")
        add_fact("quiz_limit", f"第 {k} 次测验", "minutes", minutes, [s], f"第 {k} 次测验限时多少分钟？",
                 f"第 {k} 个小测要在多长时间内做完？")

    # ---------------- 第 6 节：助教与答疑（两句一段，第二句重复名字）----------------
    new_section("助教与答疑", "【助教与答疑】本学期有 8 位助教，每人每周有一次固定的答疑时间。")
    for i, name in enumerate(TA_NAMES):          # 【Python 小知识】enumerate 同时给出下标 i（从 0 开始）和元素 name
        new_paragraph()
        add_sentence(f"助教{name}负责第 {i + 1} 次讨论课（section）。")
        slot = VALUE_POOLS["weekday_time"][int(rng.integers(len(VALUE_POOLS["weekday_time"])))]
        room = ROOMS[int(rng.integers(len(ROOMS)))]
        s = add_sentence(f"{name}的答疑时间是每{slot}，地点在{room}。")
        # 助教每年都换人，答疑安排是今年新写的 → 模型一定没见过（unknown）
        add_fact("ta_time", name, "weekday_time", slot, [s], f"助教{name}的答疑时间是什么时候？",
                 f"想当面问{name}问题，应该什么时候去？", memory="unknown")
        add_fact("ta_room", name, "room", room, [s], f"助教{name}的答疑地点在哪里？", f"去哪个房间能找到{name}？",
                 memory="unknown")

    # ---------------- 第 7 节：课程项目 ----------------
    new_section("课程项目", "【课程项目】课程项目占总评的比重最大，请尽早组队。")
    new_paragraph()
    s = add_sentence("项目以小组形式完成，每个小组最多 3 人。")
    add_fact("proj_group", "项目小组", "people", "3 人", [s], "项目小组最多几个人？", "大作业项目一组能有多少名成员？")
    s = add_sentence("每个小组有 300 美元的云计算额度，用完不再补充。")
    add_fact("proj_credit", "云计算额度", "dollars", "300 美元", [s], "每个小组有多少云计算额度？",
             "做项目能用多少钱的 GPU 服务器？")
    new_paragraph()
    pdates = [fmt_date(*_day_index_to_md(x)) for x in (40, 68, 101, 104)]   # 10/10、11/7、12/10、12/13
    s = add_sentence(f"项目提案的截止日期是 {pdates[0]}。")
    add_fact("proj_proposal", "项目提案", "date", pdates[0], [s], "项目提案的截止日期是哪天？", "项目选题方案要在哪天之前交？")
    s = add_sentence(f"中期报告的截止日期是 {pdates[1]}。")
    add_fact("proj_milestone", "中期报告", "date", pdates[1], [s], "中期报告的截止日期是哪天？", "项目中期进展报告哪天截止？")
    s = add_sentence(f"最终报告的截止日期是 {pdates[2]}，正文不超过 8 页。")
    add_fact("proj_final", "最终报告", "date", pdates[2], [s], "最终报告的截止日期是哪天？", "项目的最后一份报告什么时候交？")
    add_fact("proj_pages", "最终报告", "pages", "8 页", [s], "最终报告正文不超过几页？", "期末项目报告的篇幅限制是多少？")
    new_paragraph()
    poster_room = ROOMS[int(rng.integers(len(ROOMS)))]
    s = add_sentence(f"海报展示在 {pdates[3]} 举行，地点是{poster_room}。")
    add_fact("poster_date", "海报展示", "date", pdates[3], [s], "海报展示在哪天举行？", "项目 poster session 是哪一天？")
    add_fact("poster_room", "海报展示", "room", poster_room, [s], "海报展示的地点在哪里？", "去哪里参加项目 poster session？")

    # ---------------- 第 8 节：考试 ----------------
    new_section("考试", "【考试】本课程只有一次期中考试，没有期末考试。")
    new_paragraph()
    mid_date = fmt_date(*_day_index_to_md(64))                               # 11 月 3 日
    mid_room = ROOMS[int(rng.integers(len(ROOMS)))]
    s = add_sentence(f"期中考试在 {mid_date} 晚上 7 点举行，地点是{mid_room}。")
    add_fact("mid_date", "期中考试", "date", mid_date, [s], "期中考试在哪天举行？", "期中是几月几号考？")
    add_fact("mid_room", "期中考试", "room", mid_room, [s], "期中考试的地点在哪里？", "期中去哪个教室考？")
    s = add_sentence("考试时可以携带 2 页手写公式纸。")
    add_fact("mid_sheets", "公式纸", "sheets", "2 页", [s], "期中考试可以带几页公式纸？", "期中能带多少张小抄？")
    new_paragraph()
    s = add_sentence("需要补考的同学必须在考试前至少 5 天提交申请。")
    add_fact("mid_makeup", "补考申请", "days", "5 天", [s], "补考申请至少要提前几天提交？", "考试那天有事去不了，最晚提前多久说？")

    questions = []
    for f in facts:
        for style in ("direct", "para"):
            text = f["q_direct"] if style == "direct" else f["q_para"]
            questions.append({"qid": len(questions), "fact": f["fid"], "style": style, "text": text})
    return {"sections": sections, "sentences": sentences, "facts": facts, "questions": questions}


def _day_index_to_md(day_index):
    """
    “从 9 月 1 日起的第几天”（9 月 1 日 = 第 1 天）→ (月, 日)。

    【思路】用标准库 datetime 做日期加法，自动处理“9 月有 30 天、10 月有 31 天”这类细节。
    【Python 小知识】datetime.date(2026, 9, 1) + datetime.timedelta(days=n) 得到 n 天之后的日期。
    """
    import datetime
    d = datetime.date(2026, 9, 1) + datetime.timedelta(days=int(day_index) - 1)
    return d.month, d.day


def parse_date(text):
    """
    "10 月 23 日" → datetime.date(2026, 10, 23)。03 的 agent 做“两个日期相隔几天”时用。

    【返回】datetime.date；解析不了（比如拒答）时返回 None
    【Python 小知识】re.findall(正则表达式, s) 返回 s 里所有匹配的子串；这里的正则表达式匹配“一串连续的数字”。
    """
    import datetime
    nums = re.findall(r"\d+", text)
    if len(nums) != 2:
        return None
    return datetime.date(2026, int(nums[0]), int(nums[1]))


# =============================================================================
# 2. 分词与切块
# =============================================================================
_RUN_RE = re.compile(r"[一-鿿]+|[A-Za-z0-9]+")


def tokenize(text):
    """
    把一段文字切成检索用的“词”（token）。

    【思路】中文没有空格，正经的做法是用分词器（比如按词典切）。检索系统里一个简单、常用、效果不差的替代是
            **字的二元组（character bigram）**："截止日期" → ["截止", "止日", "日期"]。一个真正的词几乎总会
            被某个二元组完整覆盖，代价是多出一些没有意义的组合（"止日"），它们的 IDF 通常不高、影响有限。
            英文单词和数字整个作为一个 token，并转成小写。标点和空格丢掉。
    【参数】text：字符串
    【返回】token 列表（保持原来的顺序，可能有重复）
    【边界】单个汉字组成的一段（比如 "第 7 次" 里的 "第"）没有二元组，就保留这个字本身。
    【Python 小知识】re.compile(...).findall(text) 返回所有匹配的子串；U+4E00 到 U+9FFF 是常用汉字的
            Unicode 范围；列表推导式 [run[i:i + 2] for i in range(...)] 一行生成所有相邻的两个字。
    """
    tokens = []
    for run in _RUN_RE.findall(text):
        if "一" <= run[0] <= "鿿":
            if len(run) == 1:
                tokens.append(run)
            else:
                tokens.extend(run[i:i + 2] for i in range(len(run) - 1))
        else:
            tokens.append(run.lower())
    return tokens


CHUNKING_NAMES = {
    "sentence": "按句（1 句）",
    "fixed2": "固定 2 句",
    "fixed2_overlap": "固定 2 句 + 重叠 1 句",
    "paragraph": "按段落",
    "fixed6": "固定 6 句",
    "section": "整节",
}


def make_chunks(world, how="paragraph"):
    """
    把手册切成检索的单位（chunk，块）。

    【思路】检索返回的是“块”，生成时模型读的也是块。块太小，一个事实可能被切断（只有日期、不知道是哪次作业）；
            块太大，一块里塞了很多同类的事实，排序变得不准、模型读的时候也容易张冠李戴，而且上下文变长、变贵。
            六种切法（每种都只在一节之内切，不跨节）：
              "sentence"       每句一块
              "fixed2"         每 2 句一块，不管段落边界（每节开头的标题句会让后面的配对整体错开一位）
              "fixed2_overlap" 每 2 句一块，相邻两块重叠 1 句（窗口每次往后滑 1 句）
              "paragraph"      每段一块（按文档结构切：标题句自成一块）
              "fixed6"         每 6 句一块
              "section"        一节一块
    【参数】world：build_world 的返回值；how：上面六个名字之一
    【返回】块的列表，每个元素 {"cid", "section", "sids": [句子编号...], "text"}
    【边界】固定窗口切到一节末尾时，最后一块可能不满；重叠切法在一节只有 1 句时也只产生 1 块。
    """
    sents = world["sentences"]
    chunks = []

    def add(sids, sec):
        chunks.append({"cid": len(chunks), "section": sec, "sids": list(sids),
                       "text": "".join(sents[s]["text"] for s in sids)})

    for sec, section in enumerate(world["sections"]):
        all_sids = [s for para in section["paragraphs"] for s in para]
        if how == "sentence":
            for s in all_sids:
                add([s], sec)
        elif how in ("fixed2", "fixed6"):
            size = 2 if how == "fixed2" else 6
            for i in range(0, len(all_sids), size):
                add(all_sids[i:i + size], sec)
        elif how == "fixed2_overlap":
            if len(all_sids) == 1:
                add(all_sids, sec)
            for i in range(len(all_sids) - 1):
                add(all_sids[i:i + 2], sec)
        elif how == "paragraph":
            for para in section["paragraphs"]:
                add(para, sec)
        elif how == "section":
            add(all_sids, sec)
        else:
            raise ValueError(f"未知的切块方式：{how}")
    return chunks


def supports(chunk, fact):
    """
    这个块是否**完整地**包含回答该事实需要的全部句子。

    【返回】True / False
    【DL 联系】这就是本讲 recall@k 里“命中”的定义：只含日期那一句、不含“第几次作业”那一句，不算命中。
    """
    sids = set(chunk["sids"])
    return all(s in sids for s in fact["support"])


def gold_chunk(chunks, fact):
    """在一组块里找第一个完整支持 fact 的块（“标准答案”所在的块）。找不到返回 None。"""
    for c in chunks:
        if supports(c, fact):
            return c
    return None


# =============================================================================
# 3. 检索：BM25 与 TF-IDF（numpy 从零实现）
# =============================================================================
class BM25Index:
    """
    一个最小的关键词检索引擎：BM25（Robertson & Zaragoza 2009）和 TF-IDF 余弦相似度。

    【思路】先把每个块分词，建一个“块 × 词”的词频矩阵 TF（形状 (N, V)，N 个块、V 个不同的词）。
            BM25 给“块 d 里的词 t”打分：
                idf(t) · tf(t,d)·(k1 + 1) / ( tf(t,d) + k1·(1 − b + b·|d|/avgdl) )
            三个直觉：
              · idf(t) = ln(1 + (N − df + 0.5)/(df + 0.5))：越少见的词（df 小）越有区分度、分数越高；
                "截止"出现在几十块里，"林晓" 只出现在一两块里，后者更能说明“这块就是要找的”。
              · tf 的**饱和**：一个词在块里出现 1 次和 0 次差别很大，出现 10 次和 9 次差别很小（k1 控制饱和的快慢）。
              · 长度归一化：长块天然包含更多词，b 控制“按块长打折”的力度（b = 0 不打折，b = 1 完全按长度打折）。
            一个查询的分数 = 查询里每个（不同的）词在这个块上的分数之和。
            预先把每个 (块, 词) 的 BM25 权重算好存成矩阵 W，查询时只要取出查询词那几列相加。
    【参数】texts：块的文本列表；k1、b：BM25 的两个超参数（常用默认值 k1 ∈ [1.2, 2]，b = 0.75）
    【DL 联系】这是“稀疏”检索：每个块是一个 V 维、绝大多数分量为 0 的向量。讲义 §3.4 的“稠密”检索把块和查询
            都编码成几百维的嵌入向量（第 16 讲的词嵌入、第 5 讲 §6.2 的双塔编码器），用余弦相似度找最近的块。
    """

    def __init__(self, texts, k1=1.2, b=0.75):
        docs = [tokenize(t) for t in texts]
        vocab = {}
        for d in docs:
            for t in d:
                if t not in vocab:               # 字典按插入顺序记住每个词第一次出现的编号（顺序是确定的）
                    vocab[t] = len(vocab)
        N, V = len(docs), len(vocab)
        TF = np.zeros((N, V))
        for i, d in enumerate(docs):
            for t in d:
                TF[i, vocab[t]] += 1
        self.vocab, self.TF, self.k1, self.b = vocab, TF, k1, b
        self.doc_len = TF.sum(axis=1)                            # (N,) 每块的 token 数
        self.avgdl = float(self.doc_len.mean())
        self.df = (TF > 0).sum(axis=0)                           # (V,) 每个词出现在几块里
        self.idf = np.log(1.0 + (N - self.df + 0.5) / (self.df + 0.5))
        norm = k1 * (1.0 - b + b * self.doc_len / self.avgdl)    # (N,) 每块的“长度打折”
        # (N, V)：广播 —— norm[:, None] 的形状是 (N, 1)，会自动复制到 V 列；idf 的形状 (V,) 会复制到 N 行
        self.W = self.idf * TF * (k1 + 1.0) / (TF + norm[:, None])
        # TF-IDF（对照）：tf · idf'，每行归一化成单位长度，查询也一样，点积就是余弦相似度
        idf_classic = np.log(N / np.maximum(self.df, 1)) + 1.0
        tfidf = TF * idf_classic
        self.tfidf = tfidf / np.maximum(np.linalg.norm(tfidf, axis=1, keepdims=True), 1e-12)
        self.idf_classic = idf_classic

    def query_ids(self, query):
        """查询里出现在词表中的不同 token 的编号（按第一次出现的顺序）。不在词表里的词对任何块都贡献 0，直接丢掉。"""
        ids = []
        for t in tokenize(query):
            j = self.vocab.get(t)
            if j is not None and j not in ids:
                ids.append(j)
        return ids

    def scores(self, query, method="bm25"):
        """
        查询对每个块的分数。

        【参数】query：查询文字；method："bm25" 或 "tfidf"
        【返回】形状 (N,) 的分数数组
        【边界】查询里没有任何词在词表里时，所有分数都是 0（检索结果等于“按块的顺序取前 k 个”，没有意义）。
        """
        ids = self.query_ids(query)
        if not ids:
            return np.zeros(self.TF.shape[0])
        if method == "bm25":
            return self.W[:, ids].sum(axis=1)
        q = np.zeros(self.TF.shape[1])
        for t in tokenize(query):
            j = self.vocab.get(t)
            if j is not None:
                q[j] += 1.0
        q = q * self.idf_classic
        q = q / max(np.linalg.norm(q), 1e-12)
        return self.tfidf @ q

    def search(self, query, k, method="bm25"):
        """
        返回分数最高的 k 个块的编号（从高到低）。

        【返回】长度 min(k, N) 的整数列表
        【numpy 小知识】np.argsort(-s, kind="stable") 按分数从大到小排；"stable" 保证分数相同时按原来的顺序
                （块编号小的在前），这样结果每次都一样。
        """
        s = self.scores(query, method)
        order = np.argsort(-s, kind="stable")[:k]
        return [int(i) for i in order]


# =============================================================================
# 4. 模拟的 LLM
# =============================================================================
def memory_answer(fact, u_recall, u_value, params=None):
    """
    模拟 LLM 的“闭卷”回答：只靠参数记忆（预训练时见过去年的手册）。

    【思路】三种情况（概率见 PARAMS）：
              known    （今年没变）：p_recall 的概率说对，否则编一个；
              outdated （今年变了）：p_recall 的概率说出**去年的旧值**（自信地答错），否则编一个；
              unknown  （今年新增）：p_abstain_unknown 的概率说“不知道”，否则编一个看起来合理的答案。
    【参数】fact：事实字典；u_recall、u_value：两个 [0,1) 的随机数（由调用者提供，理由见 answer_question）
    【返回】(answer, outcome)：outcome ∈ {"memory", "hallucinated", "abstain"}
    【DL 联系】这对应讲义 §1.3 的“知识截止（knowledge cutoff）”和“幻觉”：模型的知识冻结在训练数据里，
            世界变了它不知道；它不知道时也常常不说不知道。
    """
    p = params or PARAMS
    if fact["memory"] in ("known", "outdated"):
        if u_recall < p["p_recall"]:
            return fact["old"], "memory"
        return hallucinate(fact["vtype"], u_value), "hallucinated"
    if u_recall < p["p_abstain_unknown"]:
        return ABSTAIN, "abstain"
    return hallucinate(fact["vtype"], u_value), "hallucinated"


def read_probability(offset, n_sentences, params=None):
    """
    信息位于上下文第 offset 句（从 0 数）、上下文共 n_sentences 句时，模型“读到并用上”它的概率。

    【思路】U 形曲线：开头和结尾最好（r_max），正中间最差；上下文越长，中间掉得越多：
                p = r_max − mid_drop · min(1, n/l_ref) · 4u(1 − u)，  u = offset / (n − 1)
            这是我们**写进模拟器的假设**，形状参考了 Liu et al.（2024，"Lost in the Middle"）在真实模型上测到的现象；
            具体数字是设计者设定的。
    【参数】offset：整数；n_sentences：整数 ≥ 1
    【返回】概率（float）
    """
    p = params or PARAMS
    if n_sentences <= 1:
        return p["r_max"]
    u = offset / (n_sentences - 1)
    scale = min(1.0, n_sentences / p["l_ref"])
    return p["r_max"] - p["mid_drop"] * scale * 4.0 * u * (1.0 - u)


def answer_question(world, fact, context, rng, prompt="plain", params=None, force_read_ok=False):
    """
    模拟 LLM 读一组资料（可能为空）后回答一个事实问题，并返回一份 trace（记录每一步发生了什么）。

    【思路】
        1. 把上下文里的块按顺序拼起来（共 n 句）。找第一个**完整支持**这个事实的块。
        2. 如果找到了（检索命中）：
             · 以 read_probability(位置, n) 的概率“读到”它（否则当作没看见，走第 3 步）；
             · 读到了，还可能张冠李戴：上下文里每多一个同类的其他事实，概率加 conf_per_distractor；
             · 没有张冠李戴，但记忆里是过时的旧值：p_conflict 的概率“相信自己的记忆”；
             · 否则答对。
        3. 没有可用的支持（检索没命中，或没读到）：
             · 上下文里有同类的值（比如别的作业的截止日期）：以 p_grab[prompt] 的概率照抄其中一个；
             · 提示词是 "grounded"（只根据资料回答，没有就说没有）：以 p_abstain_grounded 的概率拒答；
             · 否则退回到参数记忆（memory_answer）。
    【参数】
        world, fact：build_world 的世界和其中一个事实
        context     ：块的列表（检索的结果，按排名顺序）；空列表 = 闭卷
        rng         ：np.random.Generator。**函数一开始就固定地抽 9 个随机数**，后面按编号使用。这样，同一道题
                      用同一个 rng 种子在不同配置下（比如 k = 3 和 k = 5）用到的是**同一组随机数**，配置之间的差别
                      只来自配置本身，而不是随机数恰好不同（统计上叫 common random numbers，能大大减小比较的噪声）。
        prompt      ："plain"（普通提示）或 "grounded"（要求只根据资料回答）
        params      ：参数字典，默认 PARAMS
        force_read_ok：True 时，只要检索命中就一定读对（03 里“把阅读这个组件换成完美的”用）
    【返回】trace 字典：
        "answer"    ：答案字符串（或 ABSTAIN）
        "correct"   ：是否等于标准答案
        "supported" ：上下文里是否有完整支持它的块
        "offset"、"n_sent"：支持句在上下文里的位置、上下文总句数（没有支持时 offset 为 None）
        "outcome"   ："read_ok" / "confused"（张冠李戴）/ "conflict"（信了旧记忆）/ "grabbed"（照抄了别的同类值）
                       / "memory"（记忆里的值）/ "hallucinated"（编的）/ "abstain"（拒答）
        "overlooked"：检索命中了但没读到（"lost in the middle" 这一类）
    【边界】上下文里如果同一句出现两次（重叠切块），位置按第一次出现算。
    """
    p = params or PARAMS
    u = rng.random(9)                             # 固定先抽 9 个，见上面的说明（u[0]…u[8] 各管一件事）
    sents = []
    for c in context:
        sents.extend(c["sids"])
    n = len(sents)
    trace = {"answer": None, "correct": False, "supported": False, "offset": None, "n_sent": n,
             "outcome": None, "overlooked": False}

    sup = next((c for c in context if supports(c, fact)), None)
    same_kind = [f for f in world["facts"] if f["kind"] == fact["kind"]]
    visible = [f for f in same_kind if f["value_sid"] in sents]           # 上下文里能看到的同类值（含它自己）

    if sup is not None:
        trace["supported"] = True
        offset = sents.index(fact["value_sid"])
        trace["offset"] = offset
        p_read = 1.0 if force_read_ok else read_probability(offset, n, p)
        if u[0] < p_read:
            others = [f for f in visible if f["fid"] != fact["fid"]]
            p_conf = 0.0 if force_read_ok else min(p["conf_max"], p["conf_per_distractor"] * len(others))
            if others and u[1] < p_conf:
                pick = others[int(u[2] * len(others))]
                return _finish(trace, fact, pick["value"], "confused")
            if (not force_read_ok) and fact["memory"] == "outdated" and u[3] < p["p_conflict"]:
                return _finish(trace, fact, fact["old"], "conflict")
            return _finish(trace, fact, fact["value"], "read_ok")
        trace["overlooked"] = True

    # 没有可用的支持
    if visible and u[4] < p["p_grab"][prompt]:
        pick = visible[int(u[5] * len(visible))]
        return _finish(trace, fact, pick["value"], "grabbed")
    if context and prompt == "grounded" and u[6] < p["p_abstain_grounded"]:
        return _finish(trace, fact, ABSTAIN, "abstain")
    ans, outcome = memory_answer(fact, u[7], u[8], p)
    return _finish(trace, fact, ans, outcome)


def _finish(trace, fact, answer, outcome):
    """填好 trace 的答案、是否正确、结局，返回 trace。"""
    trace["answer"] = answer
    trace["correct"] = (answer == fact["value"])
    trace["outcome"] = outcome
    return trace


def rng_for(*keys):
    """
    为“某个种子、某道题、某个组件、第几次尝试”生成一个独立的随机数生成器。

    【思路】np.random.default_rng 可以接受一个整数列表当种子（内部用 SeedSequence 把它们混合）。
            同样的列表永远得到同样的随机流，不同的列表得到（统计上）独立的随机流。
            这让我们可以只换掉一个组件、而其他组件抽到的随机数完全不变（03 的误差分析要用）。
    【参数】keys：若干个非负整数
    【返回】np.random.Generator
    """
    return np.random.default_rng([int(k) for k in keys])


# =============================================================================
# 自检
# =============================================================================
if __name__ == "__main__":
    import unicodedata
    w = build_world()
    facts = w["facts"]
    print(f"手册：{len(w['sections'])} 节，{sum(len(s['paragraphs']) for s in w['sections'])} 段，"
          f"{len(w['sentences'])} 句；事实 {len(facts)} 个；问题 {len(w['questions'])} 个（直接 / 改写各一半）")
    counts = {m: sum(f["memory"] == m for f in facts) for m in ("known", "outdated", "unknown")}
    print("模型的记忆状态：" + "，".join(f"{MEMORY_NAMES[m]} {c}" for m, c in counts.items()))
    print("\n第 4 节（作业安排）的标题句和前两段：")
    sec = w["sections"][3]
    for para in sec["paragraphs"][:3]:
        print("  " + "".join(w["sentences"][s]["text"] for s in para))
    print("\n分词的例子：", tokenize("第 7 次作业的截止日期是哪天？"))
    for how in CHUNKING_NAMES:
        ch = make_chunks(w, how)
        n_sup = sum(gold_chunk(ch, f) is not None for f in facts)
        name = CHUNKING_NAMES[how]
        # 汉字和全角符号在终端里占两格（unicodedata.east_asian_width 是 "W" 或 "F"），手动补齐空格
        name += " " * (24 - sum(2 if unicodedata.east_asian_width(ch_) in ("W", "F") else 1 for ch_ in name))
        print(f"  {name}{len(ch):>4} 块；能完整支持的事实 {n_sup}/{len(facts)}")
    # 检索的健全性检查：用一句原文去查，它所在的块应该排第一
    ch = make_chunks(w, "paragraph")
    idx = BM25Index([c["text"] for c in ch])
    ok = 0
    for c in ch:
        ok += idx.search(c["text"], 1)[0] == c["cid"]
    print(f"\n用每一块的原文去检索，排第一的就是它自己：{ok}/{len(ch)}")
    # 模拟器的健全性检查：给标准答案所在的块、强制读对时，必须答对
    good = sum(answer_question(w, f, [gold_chunk(ch, f)], rng_for(0, f["fid"]), force_read_ok=True)["correct"]
               for f in facts)
    print(f"给标准答案所在的块 + 强制读对：{good}/{len(facts)} 答对")