"""
=============================================================================
CS329T 第 11 讲 · 共享模块：一个"知道真相"的玩具 RAG 世界（a simulated RAG world with known truth）
=============================================================================

对应讲义：lectures/lecture11_evaluating_llm_apps.md（§2 RAG 架构、§3 RAG triad 用到这里的每一个部件）

-----------------------------------------------------------------------------
【为什么要模拟 Why simulate?】
-----------------------------------------------------------------------------
评估 LLM 应用最难的地方是：**我们通常不知道真相**。一个回答到底对不对、是不是被检索内容支持、
检索到的段落到底相不相关 —— 现实中这些都要靠人或另一个 LLM 来判断，而判断本身也会出错。

所以本讲不调用任何真实的 LLM（不联网、不需要 API key），而是搭一个**每个事实都确切已知**的小世界：

    · 一个虚构的岛屿 Varn 上有 40 个虚构的小镇，每个镇有 6 个属性（建镇年份、人口、河流、镇长、节日、主要出口）。
      这些就是"知识库"（knowledge base）。名字都是编的，如有雷同纯属巧合。
    · 每个镇有一篇介绍文档；一部分镇还有一篇"2009 年的存档"（archive），里面的人口和镇长是**过时的**值；
      另有几篇泛泛而谈、不含具体事实的"岛屿指南"。文档被切成 chunk（chunking）。
    · 检索器（retriever）：标准库手写的 BM25、一个"语义 embedding"的玩具版（同义词归一化后的余弦相似度）、
      以及两者的 reciprocal rank fusion（hybrid）。
    · 一个模拟的 LLM 生成器（generator）：它的行为由一组**显式的旋钮**决定 ——
      上下文里有答案时照着用的概率、没有答案时说"不知道"的概率、答非所问的概率、把名字相近的两个镇搞混的概率、
      新旧信息冲突时选旧值的概率、参数记忆（预训练里"记得"的事实）的准确率、话多不多（额外 claim、客套句）、
      措辞改写的概率、"lost in the middle"（上下文中间的段落更容易被忽略）……
    · 每个回答都以**结构化的 claim**（实体, 属性, 值）保存，同时渲染成一段文本。
      于是 RAG triad 的三条边、回答是否正确，都能**精确地**算出来 —— 任何评估方法（第 11 讲的 feedback functions、
      LLM judge）都可以拿来和真相对照。

数字只说明机制，不代表任何真实 LLM 或真实 RAG 系统的性能。

-----------------------------------------------------------------------------
【接口 Interface：怎样换成真实模型】
-----------------------------------------------------------------------------
    app = RAGApp(world, AppConfig(...))
    rec = app.run(question, rng)        # → Record(q, context, resp, labels)

`generate(world, q, context, gen, rng) -> Response` 是唯一"扮演 LLM"的函数。要换成真实模型，
可以把它替换成"拼 prompt → 调用模型 → 解析回答"的函数（本课程不写任何调用外部服务的代码）。
但要注意：换成真实模型后，`true_labels` 就没有了 —— 那正是第 11 讲要解决的问题。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 rag_world.py      # 自检：打印世界的统计、一篇文档、一次完整的检索 + 生成（< 1 秒）

只用 Python 标准库（random、math、re、collections、dataclasses）。

【注释标签】
    【思路】为什么这样写   【参数】【返回】输入输出   【边界】什么时候会出问题
    【可信联系】这个函数和"可信"的哪个方面有关、哪一讲会深入
    【Python 小知识】第一次出现的写法
"""

import math
import random
import re
from collections import Counter
from dataclasses import dataclass, field
from typing import List, Optional


# =============================================================================
# 0. 小工具 Small helpers
# =============================================================================
def print_header(title):
    """打印一个醒目的分隔标题（和第 1 讲的 trust_tools.print_header 一样）。"""
    print("\n" + "=" * 100)
    print(title)
    print("=" * 100)


def mean(xs):
    """
    平均值；空列表返回 nan。

    【思路】用 math.fsum 而不是内置 sum：Python 3.12 起，内置 sum() 对浮点数用了"补偿求和"，
        结果的最后一位可能和 3.9 不同；math.fsum 在所有版本上都给出精确舍入的结果，保证输出逐字一致。
    【Python 小知识】float("nan") 是"不是一个数"（Not a Number），任何和它的比较都是 False。
    """
    xs = list(xs)
    return math.fsum(xs) / len(xs) if xs else float("nan")


# =============================================================================
# 1. 世界的原材料 Raw material of the world
# =============================================================================
# 12 个"词根"，每个生成一对名字相近的镇（confusable pair），例如 "Port Alder" 和 "Alder Falls"。
# BM25 只看词，问 "Port Alder" 时，含 "Alder" 的另一个镇的 chunk 也会被检索上来 —— 这是现实中常见的干扰。
PAIR_ROOTS = ["Alder", "Brannock", "Orrin", "Velmora", "Kestrel", "Tamsin",
              "Corve", "Halden", "Lisk", "Sorrel", "Thane", "Mirren"]
PAIR_PATTERNS = [("Port {r}", "{r} Falls"), ("Upper {r}", "{r} Bay"),
                 ("Saint {r}", "{r} Cross"), ("East {r}", "{r} Sands")]
# 16 个"独一无二"的镇名。
SINGLE_TOWNS = ["Quillhaven", "Darrowmere", "Ostrelby", "Fennickton", "Hollinvale", "Marrowgate",
                "Esterbrook", "Wrenmouth", "Calderwick", "Brisset", "Ullapane", "Gorsedown",
                "Nethercombe", "Pellingford", "Rookhurst", "Yarlow"]

ATTRS = ["founded", "population", "river", "mayor", "festival", "export"]   # 知识库里有的 6 个属性
UNANSWERABLE_ATTRS = ["area_code", "football_club", "tallest_building"]   # 知识库里**没有**的属性
STALE_ATTRS = ["population", "mayor"]                                     # 存档文档里会有过时值的属性

RIVERS = ["Kessel", "Amberline", "Tolve", "Sennet", "Corran", "Dunmere", "Halse", "Wyst",
          "Brede", "Calla", "Derrin", "Fayle", "Gorran", "Hesk", "Marle", "Oskel"]
FIRST_NAMES = ["Ilse", "Tomas", "Mira", "Odran", "Petra", "Halvard", "Wenna", "Corin",
               "Yusra", "Bastien", "Agnes", "Rafe", "Liora", "Emeric"]
LAST_NAMES = ["Damaro", "Fenwick", "Oakhart", "Selby", "Varga", "Quist", "Holloway", "Marsh",
              "Tennant", "Ruiz", "Kalder", "Brook", "Ashdown", "Pryce"]
FESTIVALS = ["Lantern Festival", "Herring Fair", "Kite Week", "Salt Moon Night", "Bell Parade",
             "Harvest Ring", "Tide Market", "Fiddlers Gathering", "Apple Wake", "Boat Blessing",
             "Wool Fair", "Torch Walk"]
EXPORTS = ["smoked herring", "slate", "wool", "cider", "glass beads", "rope", "peat", "honey",
           "oysters", "barley", "tin", "lace", "salt", "timber"]
ADJECTIVES = ["harbour", "hill", "market", "mining", "river", "fishing", "garden", "border"]
COASTS = ["north", "south", "east", "west"]
FOUNDERS = ["salt traders", "monks", "fishermen", "wool merchants", "miners", "a band of farmers"]
BUILDINGS = ["the Old Granary", "the Customs House", "St Brenna Tower", "the Clock Hall", "the Harbour Lighthouse"]
CLUB_SUFFIX = ["Rovers", "Athletic", "United", "Wanderers"]

# 问题模板。canonical 和文档用词一致；paraphrase 故意换了说法（lexical gap），只看词的 BM25 会吃亏。
QUESTION_TEMPLATES = {
    "founded": ("When was {T} founded?", "In what year did settlers establish {T}?"),
    "population": ("What is the population of {T}?", "How many residents does {T} have?"),
    "river": ("Which river runs through {T}?", "Which waterway flows through {T}?"),
    "mayor": ("Who is the mayor of {T}?", "Who leads the council of {T}?"),
    "festival": ("Which festival does {T} host every year?", "What yearly celebration is {T} known for?"),
    "export": ("What is the main export of {T}?", "What goods does {T} mainly sell abroad?"),
    "area_code": ("What is the telephone area code of {T}?", None),
    "football_club": ("Which football club plays in {T}?", None),
    "tallest_building": ("What is the tallest building in {T}?", None),
}

# 回答里的 claim 句子模板：(和文档用词一致的写法, 改写过的写法)。
CLAIM_TEMPLATES = {
    "founded": ("{T} was founded in {v}.", "{T} dates back to {v}."),
    "population": ("{T} has a population of {v}.", "Roughly {v} residents call {T} home."),
    "river": ("The {v} river runs through {T}.", "{T} sits on the banks of the {v}."),
    "mayor": ("The mayor of {T} is {v}.", "{T} is led by {v}."),
    "festival": ("{T} hosts the {v} every year.", "{T} is famous for its {v}."),
    "export": ("The main export of {T} is {v}.", "{T} mostly sells {v} to other towns."),
    "area_code": ("The area code of {T} is {v}.", "You can reach {T} on area code {v}."),
    "football_club": ("The football club of {T} is {v}.", "{v} play their home games in {T}."),
    "tallest_building": ("The tallest building in {T} is {v}.", "{v} towers over {T}."),
}

FILLERS = ["Great question!", "I hope this helps.", "Let me know if you need anything else.",
           "Varn is a lovely island to explore."]

GUIDE_DOCS = [
    "Varn visitor guide: most towns on the island were founded by traders, and many host a festival in late summer.",
    "Varn economy note: the main exports of the island are fish, slate and wool, and its rivers carry most of the trade.",
    "Varn civic guide: every town elects a mayor, and population figures are published by the island council.",
    "Varn travel tips: old river ferries connect the harbour towns, and the council runs a bus between them.",
]


# =============================================================================
# 2. 知识库 The knowledge base
# =============================================================================
@dataclass
class Fact:
    """
    一条事实（entity, attr, value），stale = True 表示它来自过时的存档文档。

    【Python 小知识】@dataclass 是一个"装饰器"：只要写出字段和类型，Python 会自动生成 __init__、__repr__ 等方法。
        于是 Fact("Port Alder", "mayor", "Ilse Damaro", False) 就能直接创建对象，f.entity 取字段。
        类型标注（: str）只是说明，Python 运行时不检查。
    """
    entity: str
    attr: str
    value: str
    stale: bool = False


class World:
    """
    整个虚构世界：镇、事实、文档。

    【思路】所有随机性都来自一个 random.Random(seed)，构造顺序固定，所以同一个 seed 永远得到同一个世界。
    【Python 小知识】random.Random(seed) 创建一个**独立**的随机数生成器，不影响全局的 random 模块；
        和 numpy 的 default_rng 一样，每个演示用自己的生成器，改一个演示不会改变另一个演示的数字。
    """

    def __init__(self, seed=11, stale_frac=0.3):
        rng = random.Random(seed)
        # ---- 镇名与"易混淆"关系 ----
        self.towns = []
        self.partner = {}                      # 镇名 → 与它名字相近的另一个镇（没有就不在字典里）
        for i, root in enumerate(PAIR_ROOTS):
            a_pat, b_pat = PAIR_PATTERNS[i % len(PAIR_PATTERNS)]
            a, b = a_pat.format(r=root), b_pat.format(r=root)
            self.towns += [a, b]
            self.partner[a], self.partner[b] = b, a
        self.towns += SINGLE_TOWNS
        # 名字里的词（小写），玩具 embedding 用它们降低"专有名词"的权重。
        self.name_tokens = sorted({w.lower() for t in self.towns for w in t.split()})

        # ---- 每个镇的真实属性 ----
        names = [f"{f} {l}" for f in FIRST_NAMES for l in LAST_NAMES]
        rng.shuffle(names)
        self.truth = {}                        # (town, attr) → 当前的真实值（字符串）
        self.extra = {}                        # 渲染文档用的修饰信息（形容词、海岸、建镇者）
        for j, t in enumerate(self.towns):
            self.truth[(t, "founded")] = str(rng.randint(1450, 1890))
            self.truth[(t, "population")] = f"{rng.randint(20, 950) * 100:,}"
            self.truth[(t, "river")] = rng.choice(RIVERS)
            self.truth[(t, "mayor")] = names[j]
            self.truth[(t, "festival")] = rng.choice(FESTIVALS)
            self.truth[(t, "export")] = rng.choice(EXPORTS)
            self.extra[t] = (rng.choice(ADJECTIVES), rng.choice(COASTS), rng.choice(FOUNDERS))

        # ---- 过时的存档：随机选 stale_frac 的镇 ----
        n_stale = int(round(stale_frac * len(self.towns)))
        self.stale_towns = sorted(rng.sample(self.towns, n_stale))
        self.stale_value = {}                  # (town, attr) → 过时的值
        k = len(self.towns)
        for t in self.stale_towns:
            cur = int(self.truth[(t, "population")].replace(",", ""))
            old = int(round(cur * rng.uniform(0.6, 0.9) / 100.0)) * 100
            self.stale_value[(t, "population")] = f"{old:,}"
            self.stale_value[(t, "mayor")] = names[k]
            k += 1

        # ---- 文档：每个文档是一串句子，每个句子记下它陈述的事实 ----
        self.docs = self._build_docs()

    # -------------------------------------------------------------------------
    def value(self, entity, attr):
        """当前的真实值；知识库里没有这个属性时返回 None。"""
        return self.truth.get((entity, attr))

    def confusables(self, entity):
        """名字相近的镇（列表，可能为空）。"""
        return [self.partner[entity]] if entity in self.partner else []

    def _build_docs(self):
        """
        生成全部文档。每个文档 = dict(doc_id, title, sents)，sents 是 [(句子文本, [Fact, ...]), ...]。

        【思路】镇的介绍文档只有第一句出现镇名，后面用 "It"、"The town" 指代 —— 这是真实文本的样子。
            于是如果切 chunk 时不给每个 chunk 加上标题（title prefix），很多 chunk 里就**没有镇名**，
            检索器不知道它属于哪个镇，读者（和 LLM）也无法把里面的事实归到某个镇上。§2 会看到这一点的代价。
        """
        docs = []
        for t in self.towns:
            adj, coast, founders = self.extra[t]
            g = self.truth
            sents = [
                (f"{t} is a small {adj} town on the {coast} coast of Varn.", []),
                (f"It was founded in {g[(t, 'founded')]} by {founders}.", [Fact(t, "founded", g[(t, "founded")])]),
                (f"The town has a population of {g[(t, 'population')]}.", [Fact(t, "population", g[(t, "population")])]),
                (f"The {g[(t, 'river')]} river runs through the town centre.", [Fact(t, "river", g[(t, "river")])]),
                (f"The current mayor is {g[(t, 'mayor')]}.", [Fact(t, "mayor", g[(t, "mayor")])]),
                (f"Every year it hosts the {g[(t, 'festival')]}.", [Fact(t, "festival", g[(t, "festival")])]),
                (f"Its main export is {g[(t, 'export')]}.", [Fact(t, "export", g[(t, "export")])]),
            ]
            docs.append(dict(doc_id=f"town:{t}", title=f"{t} (town profile)", sents=sents))
        for t in self.stale_towns:
            sv = self.stale_value
            sents = [
                (f"Archive note, 2009: {t} then had a population of {sv[(t, 'population')]}.",
                 [Fact(t, "population", sv[(t, "population")], stale=True)]),
                (f"In 2009 the mayor of {t} was {sv[(t, 'mayor')]}.",
                 [Fact(t, "mayor", sv[(t, "mayor")], stale=True)]),
            ]
            docs.append(dict(doc_id=f"archive:{t}", title=f"{t} (archive, 2009)", sents=sents))
        for i, text in enumerate(GUIDE_DOCS):
            docs.append(dict(doc_id=f"guide:{i}", title="Varn guide", sents=[(text, [])]))
        return docs


# =============================================================================
# 3. 文本处理与 chunking Text processing and chunking
# =============================================================================
STOPWORDS = {
    "a", "an", "the", "of", "in", "on", "is", "was", "were", "what", "which", "who", "when", "does", "do",
    "did", "to", "and", "by", "it", "its", "there", "that", "this", "for", "at", "with", "from", "as", "be",
    "are", "has", "have", "had", "i", "you", "me", "my", "your", "how", "many", "much", "then", "today",
    "answer", "according", "records", "could", "not", "find", "provided", "documents", "say", "s",
}


def tokenize(text):
    """
    把文本切成小写的词和数字。"18,400" 先去掉千分位逗号变成 "18400"，免得被切成两个词。

    【Python 小知识】re.sub(pattern, repl, s) 用正则表达式替换；(?<=\\d),(?=\\d) 的意思是
        "前面是数字、后面也是数字的逗号"（(?<=...) 和 (?=...) 是"环视"，只检查、不吃掉字符）。
        re.findall(r"[a-z0-9]+", s) 找出所有由小写字母和数字组成的连续片段。
    """
    text = re.sub(r"(?<=\d),(?=\d)", "", text.lower())
    return re.findall(r"[a-z0-9]+", text)


def content_tokens(text):
    """去掉停用词之后的词。BM25、词重叠类的打分都用它。"""
    return [w for w in tokenize(text) if w not in STOPWORDS]


@dataclass
class Chunk:
    """
    一个 chunk（检索的基本单位）。

    facts：这个 chunk 里**能归属到某个镇**的事实（chunk 文本里出现了那个镇名）；
    hidden：这个 chunk 里陈述了、但因为没有镇名而**无法归属**的事实（例如 "It was founded in 1642."）。
    """
    cid: int
    doc_id: str
    text: str
    n_sents: int
    facts: List[Fact] = field(default_factory=list)
    hidden: List[Fact] = field(default_factory=list)


def build_chunks(world, chunk_sents=2, title_prefix=True):
    """
    把文档切成 chunk。

    【思路】按句子做**不重叠**的窗口，每 chunk_sents 句一个 chunk。title_prefix = True 时，
        在每个 chunk 前面加上 "Port Alder (town profile): " 这样的标题（一种常见的 "contextual chunk header" 做法）。
        一个事实能否"归属"，取决于 chunk 文本里有没有它的实体名 —— 这是模拟"读者能否看出这句话在说哪个镇"。
    【参数】world；chunk_sents：每个 chunk 几句（1 = 句子级，7 = 整篇文档）；title_prefix：是否加标题
    【返回】Chunk 列表，cid 从 0 开始连续编号
    【边界】chunk_sents ≥ 文档长度时整篇文档就是一个 chunk；指南文档没有事实，facts 为空。
    【可信联系】chunk 太小 → 指代丢失（"It" 是谁？），检索不到、也无法 grounding；chunk 太大 → 上下文里塞满无关内容。
        这是 RAG 的第一个设计选择，也是 context relevance 这条边最常见的失败原因。
    """
    chunks = []
    for d in world.docs:
        sents = d["sents"]
        for s in range(0, len(sents), chunk_sents):
            window = sents[s:s + chunk_sents]
            body = " ".join(txt for txt, _ in window)
            text = f"{d['title']}: {body}" if title_prefix else body
            ch = Chunk(cid=len(chunks), doc_id=d["doc_id"], text=text, n_sents=len(window))
            for _, facts in window:
                for f in facts:
                    (ch.facts if f.entity in text else ch.hidden).append(f)
            chunks.append(ch)
    return chunks


# =============================================================================
# 4. 检索器 Retrievers
# =============================================================================
class BM25:
    """
    Okapi BM25（Robertson & Zaragoza 2009 的综述）：经典的"词袋"检索打分，至今仍是很强的基线。

        score(q, d) = Σ_{t ∈ q} IDF(t) · tf(t, d)·(k1 + 1) / ( tf(t, d) + k1·(1 − b + b·|d|/avgdl) )
        IDF(t) = log(1 + (N − n_t + 0.5)/(n_t + 0.5))

    【思路】三个直觉：(1) 罕见的词更有信息量（IDF）；(2) 一个词出现多次有帮助，但收益递减（k1 控制饱和速度）；
        (3) 长文档天然含更多词，要按长度打折（b 控制打折力度）。k1 = 1.5、b = 0.75 是常用默认值。
    【边界】查询里的词如果一个文档都没有出现，所有分数都是 0，排序退化为按 cid —— 这正是 paraphrase 问题的失败方式。
    """

    def __init__(self, texts, k1=1.5, b=0.75):
        self.k1, self.b = k1, b
        toks = [content_tokens(t) for t in texts]
        self.tf = [Counter(t) for t in toks]
        self.dl = [len(t) for t in toks]
        self.N = len(texts)
        self.avgdl = math.fsum(self.dl) / max(self.N, 1)
        df = Counter()
        for t in toks:
            df.update(set(t))           # 每个文档对每个词只计一次
        self.idf = {w: math.log(1.0 + (self.N - n + 0.5) / (n + 0.5)) for w, n in df.items()}

    def scores(self, query):
        """对每个文档返回 BM25 分数（列表，顺序和构造时的 texts 一样）。"""
        q = content_tokens(query)
        out = []
        for i in range(self.N):
            s = 0.0
            norm = self.k1 * (1.0 - self.b + self.b * self.dl[i] / self.avgdl)
            for w in q:                  # 按查询词的固定顺序累加，保证浮点结果在各 Python 版本上一致
                f = self.tf[i].get(w, 0)
                if f:
                    s += self.idf[w] * f * (self.k1 + 1.0) / (f + norm)
            out.append(s)
        return out


# 玩具"语义 embedding"用的同义词表：把不同说法映射到同一个"概念"。
# 真实的 dense embedding 是从数据里学到这种等价关系的；这里手写，好让机制一目了然。
SYNONYMS = {
    "founded": "found", "establish": "found", "established": "found", "settlers": "found",
    "residents": "population", "people": "population", "inhabitants": "population",
    "waterway": "river", "flows": "run", "runs": "run",
    "leads": "mayor", "council": "mayor", "led": "mayor",
    "celebration": "festival", "yearly": "annual", "every": "annual", "year": "annual", "hosts": "host",
    "goods": "export", "sell": "export", "sells": "export", "abroad": "export", "exports": "export",
}


def embed(text, name_tokens, name_weight=0.5):
    """
    玩具 dense embedding：同义词归一化后的"概念袋"（bag of concepts），专有名词降权。

    【思路】dense retriever 擅长"意思相同、用词不同"（paraphrase），但对罕见的专有名词往往不如 BM25 精确
        —— 名字在 embedding 空间里会和长得像的名字挤在一起。我们用两个简单的规则模拟这两点：
        (1) SYNONYMS 把 "residents" 和 "population" 映射成同一个维度；(2) 镇名里的词权重只有 name_weight。
    【参数】text；name_tokens：所有镇名里的词（小写）；name_weight：专有名词的权重
    【返回】dict：概念 → 权重
    【可信联系】这是"检索质量"的一个可控旋钮：BM25 vs dense vs hybrid 在不同问题切片上谁更好，§2 用数字看。
    """
    names = set(name_tokens)
    vec = Counter()
    for w in content_tokens(text):
        c = SYNONYMS.get(w, w)
        vec[c] += name_weight if w in names else 1.0
    return dict(vec)


def cosine(u, v):
    """
    两个稀疏向量（dict）的余弦相似度。

    【边界】任何一个是零向量时返回 0。
    【Python 小知识】sorted(set(u) & set(v)) 取两个字典共同的键并排序 —— 排序是为了让浮点累加顺序固定。
    """
    common = sorted(set(u) & set(v))
    dot = math.fsum(u[k] * v[k] for k in common)
    nu = math.sqrt(math.fsum(x * x for x in u.values()))
    nv = math.sqrt(math.fsum(x * x for x in v.values()))
    return dot / (nu * nv) if nu > 0 and nv > 0 else 0.0


def ranking_from_scores(scores):
    """按分数从高到低排序的下标；分数相同时按下标从小到大（保证确定性）。"""
    return sorted(range(len(scores)), key=lambda i: (-scores[i], i))


class Retriever:
    """
    三种检索器的统一接口：kind ∈ {"bm25", "dense", "hybrid"}。

    【思路】hybrid 用 reciprocal rank fusion（RRF，Cormack et al. 2009）：score(d) = Σ_r 1/(60 + rank_r(d))，
        把两个排序"投票"合并，不需要把两种分数放到同一个尺度上。
        full_ranking 会缓存（同一个问题只算一次），top-k 只是截取前 k 个。
    【可信联系】检索器决定了 RAG triad 里 context relevance 这条边；它的失败会一路传到 groundedness 和 answer relevance。
    """

    def __init__(self, chunks, kind, name_tokens):
        self.chunks, self.kind = chunks, kind
        self.name_tokens = name_tokens
        texts = [c.text for c in chunks]
        self.bm25 = BM25(texts)
        self.vecs = [embed(t, name_tokens) for t in texts]
        self._cache = {}

    def full_ranking(self, query):
        if query in self._cache:
            return self._cache[query]
        if self.kind == "bm25":
            rank = ranking_from_scores(self.bm25.scores(query))
        elif self.kind == "dense":
            qv = embed(query, self.name_tokens)
            rank = ranking_from_scores([cosine(qv, v) for v in self.vecs])
        elif self.kind == "hybrid":
            r1 = ranking_from_scores(self.bm25.scores(query))
            qv = embed(query, self.name_tokens)
            r2 = ranking_from_scores([cosine(qv, v) for v in self.vecs])
            fused = [0.0] * len(self.chunks)
            for r in (r1, r2):
                for pos, i in enumerate(r):
                    fused[i] += 1.0 / (60.0 + pos + 1)
            rank = ranking_from_scores(fused)
        else:
            raise ValueError(f"unknown retriever kind: {self.kind}")
        self._cache[query] = rank
        return rank

    def retrieve(self, query, k):
        """返回前 k 个 Chunk（按排名）。"""
        return [self.chunks[i] for i in self.full_ranking(query)[:k]]


# =============================================================================
# 5. 问题 Questions
# =============================================================================
@dataclass
class Question:
    """
    一个评估问题。slice（切片）按优先级 unanswerable > stale > paraphrase > confusable > simple 指定：

        unanswerable：知识库里没有的属性（正确行为是说"不知道"）
        stale       ：这个（镇, 属性）在存档里有一个过时的值，上下文里可能新旧并存
        paraphrase  ：问题换了说法，和文档没有共同的关键词
        confusable  ：这个镇有一个名字相近的"邻居"
        simple      ：以上都不是
    """
    qid: int
    text: str
    entity: str
    attr: str
    template: str
    slice: str
    answerable: bool
    gold: Optional[str]


SLICES = ["simple", "confusable", "paraphrase", "stale", "unanswerable"]


def make_questions(world):
    """
    生成全部问题（顺序固定）：40 个镇 × 6 个属性 × 2 种问法 + 40 个镇 × 3 个无法回答的属性 = 600 题。

    【返回】Question 列表
    【参数】world：World
    【边界】paraphrase 模板只对知识库里有的 6 个属性存在；无法回答的属性只有一种问法。
    【可信联系】切片（slice）是评估集设计的核心：总体准确率会掩盖某一类问题上的系统性失败（第 1 讲 §5.2 的"按群体拆开"）。
    """
    qs = []
    for t in world.towns:
        for a in ATTRS:
            for ti, tmpl in enumerate(QUESTION_TEMPLATES[a]):
                template = "canonical" if ti == 0 else "paraphrase"
                if (t, a) in world.stale_value:
                    sl = "stale"
                elif template == "paraphrase":
                    sl = "paraphrase"
                elif t in world.partner:
                    sl = "confusable"
                else:
                    sl = "simple"
                qs.append(Question(len(qs), tmpl.format(T=t), t, a, template, sl, True, world.value(t, a)))
        for a in UNANSWERABLE_ATTRS:
            qs.append(Question(len(qs), QUESTION_TEMPLATES[a][0].format(T=t), t, a, "canonical",
                               "unanswerable", False, None))
    return qs


# =============================================================================
# 6. 模拟的 LLM 生成器 The simulated LLM
# =============================================================================
@dataclass
class GenConfig:
    """
    模拟 LLM 的全部旋钮（knobs）。每个旋钮对应文献里报告过的一种行为。

    model         ：底座模型的名字。参数记忆（memory）由它决定 —— 同一个底座模型对同一个事实总是"记得"同一个值。
    p_follow      ：上下文里有答案时，照着上下文回答的概率（否则用参数记忆）。prompt 里写不写"只根据上下文回答"主要改它。
    p_abstain     ：上下文里没有答案时，说"不知道"的概率（否则凭记忆编一个 —— 也就是 hallucination）。
    p_off_topic   ：答非所问的概率：回答的是同一个镇的另一个属性（比如 query rewriting 出错、模板 bug）。
    p_confuse     ：上下文里没有本镇的答案、但有名字相近的镇的同类事实时，张冠李戴的概率；
                    上下文里有本镇答案时，这个概率除以 3。
    p_pick_stale  ：上下文里新旧两个值并存时，选旧值的概率（knowledge conflict）。
    p_know        ：参数记忆里恰好记得正确值的概率（世界是虚构的，所以默认很低）。
    n_extra       ：额外附带的 claim 条数（verbosity）。
    p_extra_memory：额外 claim 来自记忆（未必有依据）而不是来自上下文的概率。
    n_filler      ：客套句条数（"Great question!"）—— 不需要核实，但会干扰一些 groundedness 打分。
    p_paraphrase  ：每个 claim 用改写措辞的概率（影响基于词重叠的打分）。
    style         ：回答风格（"A" 或 "B"），LLM judge 的 self-preference 偏差会用到它。
    middle_penalty：lost in the middle（Liu et al. 2023）：k ≥ 3 时，排在中间位置的段落被采用的概率乘以 (1 − middle_penalty)。
    long_penalty  ：长 chunk 的干扰（Shi et al. 2023, "LLMs can be easily distracted by irrelevant context"）：
                    答案所在的 chunk 每比 2 句多一句，被采用的概率乘以 (1 − long_penalty)。
    leaked        ：数据污染（contamination）：预训练语料里混进了某个公开评估集的答案。
                    这里列出的（镇, 属性）在参数记忆里以 0.95 的概率是正确值（其余事实仍然是 p_know）。
    """
    model: str = "llm-A"
    p_follow: float = 0.95
    p_abstain: float = 0.85
    p_off_topic: float = 0.02
    p_confuse: float = 0.3
    p_pick_stale: float = 0.35
    p_know: float = 0.15
    n_extra: int = 0
    p_extra_memory: float = 0.5
    n_filler: int = 0
    p_paraphrase: float = 0.2
    style: str = "A"
    middle_penalty: float = 0.15
    long_penalty: float = 0.03
    leaked: frozenset = frozenset()


@dataclass
class Claim:
    """回答里的一条 claim：(entity, attr, value) + 来源（context / memory / confused）+ 渲染出的句子。"""
    entity: str
    attr: str
    value: str
    origin: str
    paraphrased: bool = False
    text: str = ""


@dataclass
class Response:
    """
    一个回答：结构化的 claims（main 是回答问题的那条）+ 客套句 + 渲染出的文本。

    abstained = True 表示模型说了"不知道"，此时 claims 为空、main 为 None。
    """
    claims: List[Claim]
    main: Optional[Claim]
    abstained: bool
    fillers: List[str]
    text: str
    style: str
    n_extra: int


def fake_value(attr, entity, truth, r):
    """
    给一个属性编一个"看起来合理"的值（和真实值不同）。参数记忆出错、或无法回答的问题被硬答时用它。

    【参数】attr；entity：镇名（编 football club 名字时要用）；truth：真实值（None 表示知识库里没有）；r：随机数生成器
    【返回】字符串
    """
    for _ in range(50):
        if attr == "founded":
            v = str(r.randint(1450, 1890))
        elif attr == "population":
            v = f"{r.randint(20, 950) * 100:,}"
        elif attr == "river":
            v = r.choice(RIVERS)
        elif attr == "mayor":
            v = f"{r.choice(FIRST_NAMES)} {r.choice(LAST_NAMES)}"
        elif attr == "festival":
            v = r.choice(FESTIVALS)
        elif attr == "export":
            v = r.choice(EXPORTS)
        elif attr == "area_code":
            v = f"0{r.randint(1000, 9999)}"
        elif attr == "football_club":
            v = f"{entity} {r.choice(CLUB_SUFFIX)}"
        else:
            v = r.choice(BUILDINGS)
        if v != truth:
            return v
    return v


def memory_value(world, model, entity, attr, p_know, leaked=False):
    """
    模型的"参数记忆"：对某个事实，它记得的值。

    【思路】同一个底座模型对同一个事实的记忆应该是**固定的**（不是每次重新抽签），所以用
        random.Random(f"mem|{model}|{entity}|{attr}") 这个只依赖于（模型, 事实）的生成器：
        先抽一个数决定"记不记得"，记不得就编一个看似合理的值。
    【Python 小知识】用字符串做 random.Random 的种子时，Python 用 SHA-512 把它变成整数，
        结果不受 PYTHONHASHSEED 影响、在各个 Python 版本上都一样 —— 这是"确定性的伪随机"的一个方便写法。
    【可信联系】"有时 hallucination 恰好是对的"（Fall 2023 第 2 讲的 slide）：模型凭记忆答对了，
        但回答没有依据（ungrounded），你也无法验证。p_know 控制这种情况有多常见。
        leaked = True 模拟数据污染：这个事实出现在了预训练语料里（例如某个公开评估集连同答案被爬进了训练数据）。
    """
    r = random.Random(f"mem|{model}|{entity}|{attr}")
    truth = world.value(entity, attr)
    if truth is not None and r.random() < (0.95 if leaked else p_know):
        return truth
    return fake_value(attr, entity, truth, r)


def render_claim(c, rng, p_paraphrase):
    """把一条 claim 渲染成句子（以 p_paraphrase 的概率用改写过的模板）。原地写入 c.text、c.paraphrased。"""
    canon, para = CLAIM_TEMPLATES[c.attr]
    c.paraphrased = rng.random() < p_paraphrase
    c.text = (para if c.paraphrased else canon).format(T=c.entity, v=c.value)
    return c


def generate(world, q, context, gen, rng):
    """
    模拟的 LLM：读问题和上下文，写一个回答。

    【思路】一步一步模拟一个"大体靠谱、但会以已知方式犯错"的模型：
        1. 以 p_off_topic 的概率把要回答的属性换成别的（答非所问）。
        2. 在上下文里找"本镇 + 该属性"的事实（只算能归属的，即 chunk 里有镇名的）。
           · 找到了：新旧并存时以 p_pick_stale 选旧值；以 p_confuse/3 的概率被名字相近的镇带偏；
             否则以 p_follow ×（位置系数）×（长 chunk 系数）的概率照着用，剩下的概率用参数记忆。
           · 没找到：有"邻居镇"的同类事实时，以 p_confuse 的概率张冠李戴；否则以 p_abstain 的概率说不知道，
             剩下的概率凭参数记忆回答（hallucination）。
        3. 附加 n_extra 条额外 claim（verbosity）：以 p_extra_memory 的概率来自记忆，否则从上下文里摘。
        4. 加 n_filler 句客套话，按 style 渲染成文本。
    【参数】world；q：Question；context：Chunk 列表（按检索排名）；gen：GenConfig；rng：random.Random
    【返回】Response
    【边界】上下文为空时只能弃权或凭记忆回答；无法回答的问题在上下文里永远找不到事实。
    【可信联系】每个分支都对应 RAG 的一个失败模式（§2 的表）；triad 的三条边各自"看见"其中的一部分（§3）。
    """
    target = q.attr
    if rng.random() < gen.p_off_topic:
        others = [a for a in ATTRS if a != q.attr]
        target = others[rng.randrange(len(others))]

    # 上下文里的候选事实，记下它们所在 chunk 的位置（0 = 排名第一）。
    own, near = [], []
    neighbours = world.confusables(q.entity)
    for pos, ch in enumerate(context):
        for f in ch.facts:
            if f.attr != target:
                continue
            if f.entity == q.entity:
                own.append((pos, f))
            elif f.entity in neighbours:
                near.append((pos, f))

    def from_memory(attr):
        v = memory_value(world, gen.model, q.entity, attr, gen.p_know, (q.entity, attr) in gen.leaked)
        return Claim(q.entity, attr, v, "memory")

    main, abstained = None, False
    if own:
        current = [pf for pf in own if not pf[1].stale]
        stale = [pf for pf in own if pf[1].stale]
        if current and stale:
            pos, fact = stale[0] if rng.random() < gen.p_pick_stale else current[0]
        else:
            pos, fact = own[0]
        middle = len(context) >= 3 and 0 < pos < len(context) - 1
        p_use = gen.p_follow * ((1.0 - gen.middle_penalty) if middle else 1.0)
        p_use *= (1.0 - gen.long_penalty) ** max(0, context[pos].n_sents - 2)
        if near and rng.random() < gen.p_confuse / 3.0:
            main = Claim(q.entity, target, near[0][1].value, "confused")
        elif rng.random() < p_use:
            main = Claim(q.entity, target, fact.value, "context")
        else:
            main = from_memory(target)
    else:
        if near and rng.random() < gen.p_confuse:
            main = Claim(q.entity, target, near[0][1].value, "confused")
        elif rng.random() < gen.p_abstain:
            abstained = True
        else:
            main = from_memory(target)

    claims = [main] if main is not None else []
    n_extra = 0
    if not abstained:
        ctx_facts = [f for ch in context for f in ch.facts]
        for _ in range(gen.n_extra):
            used = {(c.entity, c.attr) for c in claims}
            if ctx_facts and rng.random() >= gen.p_extra_memory:
                pool = [f for f in ctx_facts if (f.entity, f.attr) not in used]
                if not pool:
                    continue
                f = pool[rng.randrange(len(pool))]
                claims.append(Claim(f.entity, f.attr, f.value, "context"))
            else:
                pool = [a for a in ATTRS if (q.entity, a) not in used]
                if not pool:
                    continue
                claims.append(from_memory(pool[rng.randrange(len(pool))]))
            n_extra += 1
    for c in claims:
        render_claim(c, rng, gen.p_paraphrase)

    fillers = [FILLERS[rng.randrange(len(FILLERS))] for _ in range(gen.n_filler)] if not abstained else []
    prefix = "Answer: " if gen.style == "A" else "According to the records: "
    if abstained:
        body = "I could not find this in the provided documents." if gen.style == "A" else "the documents do not say."
    else:
        body = " ".join([c.text for c in claims] + fillers)
    return Response(claims, main, abstained, fillers, prefix + body, gen.style, n_extra)


# =============================================================================
# 7. 真相：RAG triad 与正确性 Ground truth for the triad and for correctness
# =============================================================================
def chunk_is_relevant(ch, q):
    """chunk 里有没有"本镇 + 该属性"的事实（新旧都算，只算能归属的）。"""
    return any(f.entity == q.entity and f.attr == q.attr for f in ch.facts)


def claim_supported(c, context):
    """一条 claim 是否被上下文里某个 chunk 支持：(entity, attr, value) 完全一致，而且该事实在 chunk 里能归属。"""
    return any(f.entity == c.entity and f.attr == c.attr and f.value == c.value
               for ch in context for f in ch.facts)


def true_labels(world, q, context, resp):
    """
    用结构化的真相精确计算 RAG triad 和正确性。

    【思路】定义（§3.1 会正式写出公式）：
        context relevance (CR)：上下文里"相关" chunk 的比例（chunk 级，和 TruLens 的做法一致）；
        groundedness (G)      ：回答里被上下文支持的 claim 的比例（没有 claim —— 弃权 —— 记为 1）；
        answer relevance (AR) ：主 claim 是否在回答被问的那个（镇, 属性）。弃权记为 0（它没有回答问题）；
        correct               ：可回答的问题：主 claim 的值等于**当前**真实值；无法回答的问题：弃权才算对。
        还额外记录：answer_in_context（当前正确值在不在上下文里，= 检索的 recall）。
    【参数】world；q：Question；context：Chunk 列表；resp：Response
    【返回】dict
    【边界】AR 只看"是不是在回答这个问题"，不看对不对 —— 一个切题但错误的回答 AR = 1。
    """
    rel = [1.0 if chunk_is_relevant(ch, q) else 0.0 for ch in context]
    cr = mean(rel) if rel else 0.0
    sup = [claim_supported(c, context) for c in resp.claims]
    g = mean([1.0 if s else 0.0 for s in sup]) if sup else 1.0
    m = resp.main
    on_topic = m is not None and m.entity == q.entity and m.attr == q.attr
    ar = 1.0 if on_topic else 0.0
    if q.answerable:
        correct = on_topic and m.value == q.gold
    else:
        correct = resp.abstained
    ans_in_ctx = q.answerable and any(f.entity == q.entity and f.attr == q.attr and not f.stale
                                      for ch in context for f in ch.facts)
    return dict(cr=cr, g=g, ar=ar, correct=1 if correct else 0, ans_in_ctx=1 if ans_in_ctx else 0,
                fully_grounded=1 if all(sup) else 0, abstained=1 if resp.abstained else 0,
                main_origin=(m.origin if m is not None else "abstain"))


# =============================================================================
# 8. 一个完整的 RAG 应用 The RAG app
# =============================================================================
@dataclass
class AppConfig:
    """
    一个 RAG 应用版本 = 检索配置 + 生成配置。

    【Python 小知识】field(default_factory=GenConfig) 让每个 AppConfig 得到**自己的**一个 GenConfig 对象；
        如果写成 gen: GenConfig = GenConfig()，所有实例会共享同一个默认对象（dataclass 也不允许这样写可变默认值）。
    """
    name: str
    retriever: str = "bm25"
    k: int = 3
    chunk_sents: int = 2
    title_prefix: bool = True
    gen: GenConfig = field(default_factory=GenConfig)


@dataclass
class Record:
    """一次运行的完整记录（相当于 TruLens 里的一条 trace）：问题、检索到的上下文、回答、真相标签。"""
    app: str
    q: Question
    context: List[Chunk]
    resp: Response
    labels: dict


_INDEX_CACHE = {}


def get_retriever(world, chunk_sents, title_prefix, kind):
    """
    构建（并缓存）一个检索器。同一个 world + 配置只建一次索引。

    【Python 小知识】模块级的字典 _INDEX_CACHE 充当缓存；键是一个元组（tuple），元组可以做字典的键，列表不行。
        id(world) 是对象在内存中的编号，用来区分不同的 World 对象。
    """
    key = (id(world), chunk_sents, title_prefix, kind)
    if key not in _INDEX_CACHE:
        chunks = build_chunks(world, chunk_sents, title_prefix)
        _INDEX_CACHE[key] = Retriever(chunks, kind, world.name_tokens)
    return _INDEX_CACHE[key]


class RAGApp:
    """
    检索 → 拼上下文 → 生成。

    【思路】run() 返回一条 Record。生成器的随机性来自调用者传入的 rng，所以每个演示可以用自己的生成器。
    """

    def __init__(self, world, cfg):
        self.world, self.cfg = world, cfg
        self.retriever = get_retriever(world, cfg.chunk_sents, cfg.title_prefix, cfg.retriever)

    def run(self, q, rng):
        ctx = self.retriever.retrieve(q.text, self.cfg.k)
        resp = generate(self.world, q, ctx, self.cfg.gen, rng)
        return Record(self.cfg.name, q, ctx, resp, true_labels(self.world, q, ctx, resp))


# =============================================================================
# 自检 Self-check
# =============================================================================
if __name__ == "__main__":
    w = World()
    qs = make_questions(w)
    print_header("rag_world.py 自检：世界的统计")
    print(f"镇：{len(w.towns)}（{len(w.partner)} 个有名字相近的邻居），存档镇：{len(w.stale_towns)}，文档：{len(w.docs)}")
    cnt = Counter(q.slice for q in qs)
    print("问题数：", len(qs), "；按切片：", ", ".join(f"{s}={cnt[s]}" for s in SLICES))
    for cs, tp in [(2, True), (2, False), (1, True), (7, True)]:
        ch = build_chunks(w, cs, tp)
        n_hidden = sum(len(c.hidden) for c in ch)
        print(f"chunk_sents={cs}, title_prefix={tp}: {len(ch)} 个 chunk，无法归属的事实 {n_hidden} 条")
    t = w.towns[0]
    print("\n一篇文档：", " ".join(s for s, _ in w.docs[0]["sents"]))
    app = RAGApp(w, AppConfig("demo"))
    rng = random.Random(0)
    q = qs[0]
    rec = app.run(q, rng)
    print("\n问题：", q.text, "| 真实答案：", q.gold)
    for i, ch in enumerate(rec.context):
        print(f"  context[{i}] (cid {ch.cid}): {ch.text}")
    print("回答：", rec.resp.text)
    print("真相标签：", rec.labels)