"""
=============================================================================
第 1 讲 · 工具模块：BM25 检索器
=============================================================================

对应讲义：lectures/lecture01_intro_hallucination.md 第 4.1 节（RAG 的流程）

-----------------------------------------------------------------------------
【这个模块做什么？】
-----------------------------------------------------------------------------
检索增强生成（RAG）的第一步是"检索"：给一个查询（query），从一大堆段落里找出最相关的 k 段。
这里实现最经典的稀疏检索算法 BM25（Robertson 等，1994 年左右定型，至今仍是很强的基线）。
第 5 讲会讲稠密检索（用神经网络把文字变成向量）和 ColBERT，本讲只需要 BM25。

BM25 的直觉（一个词对"查询—段落"匹配分数的贡献）：
    1. 查询里的词在段落里出现得越多，越相关 —— 但收益递减（出现 10 次不比出现 3 次相关多少）。
       这由参数 k1 控制。
    2. 越罕见的词越重要：所有段落都有的 "the" 不提供信息，只有一篇文章里出现的 "velmora" 信息量很大。
       这就是 IDF（逆文档频率）。
    3. 长段落天然包含更多的词，要打个折扣，否则长段落总是占便宜。这由参数 b 控制。

公式（对查询里的每个词 t 求和）：
                         tf(t, P) · (k1 + 1)
    score(Q, P) = Σ IDF(t) · ─────────────────────────────────────
                  t∈Q        tf(t, P) + k1 · (1 − b + b · |P| / avgdl)

    tf(t, P)：词 t 在段落 P 里出现的次数；|P|：段落的词数；avgdl：所有段落的平均词数
    IDF(t) = ln( (N − n_t + 0.5) / (n_t + 0.5) + 1 )，N：段落总数，n_t：含有 t 的段落数

【它的弱点】BM25 只认"字面上相同的词"：问题里写 "genetic cause"，文章里写 "caused by mutations
in the ... gene"，它就匹配不上（cause ≠ caused，genetic ≠ gene）。这正是 02 号文件要演示的
"检索没有返回答案"的来源之一。

-----------------------------------------------------------------------------
【用法】
-----------------------------------------------------------------------------
    from retriever import BM25
    bm25 = BM25(world.passages)
    top = bm25.search("Which gene is mutated in Velmora syndrome?", k=3)   # list[Passage]
"""

import math
import re
from collections import Counter

# 停用词：太常见、几乎不提供信息的词。去掉它们能让分数更集中在有意义的词上。
# 【边界】这个表是为本讲的英文玩具语料手工挑的，不是通用的停用词表。
STOPWORDS = {
    "a", "an", "the", "of", "in", "is", "are", "was", "were", "by", "and", "or", "to", "for", "on",
    "what", "which", "who", "when", "how", "does", "do", "did", "about", "it", "its", "that", "this",
    "with", "be", "me", "tell", "md",
}


def tokenize(text):
    """
    把文字切成词（小写、只保留字母和数字、去掉停用词）。

    【参数】text：str
    【返回】list[str]，例如 tokenize("Which gene is mutated in Velmora syndrome?")
            → ['gene', 'mutated', 'velmora', 'syndrome']
    【边界】
        - 不做词干还原（stemming）：caused 和 cause 被当成两个不同的词。真实系统常用
          Porter 词干器把它们都变成 "caus"，这里故意不做，好让"字面不匹配"的问题更明显。
        - "1 in 50,000" 会被切成 ['1', '50', '000']。

    【Python 小知识】re.findall(r"[a-z0-9]+", s) 找出 s 里所有"连续的字母或数字"片段。
    """
    return [w for w in re.findall(r"[a-z0-9]+", text.lower()) if w not in STOPWORDS]


class BM25:
    """
    BM25 检索器。

    【构造参数】
        passages : list[Passage]（toy_world.Passage；只用到 .text 和 .title）
        k1       : 词频饱和参数，常用 1.2~2.0。越大，"出现多次"的奖励越大
        b        : 长度归一化参数，0~1。0 = 不考虑段落长度，1 = 完全按长度归一化
    【边界】段落列表在构造时就固定下来；想加入新文档，需要重新构造（这里不实现增量索引）。
    """

    def __init__(self, passages, k1=1.5, b=0.75):
        self.passages = list(passages)
        self.k1, self.b = k1, b
        # 每段的词频表。标题也算进段落的文字里（很多真实系统也这么做：标题往往是最重要的线索）。
        self.tfs = [Counter(tokenize(p.title + " " + p.text)) for p in self.passages]
        self.lengths = [sum(tf.values()) for tf in self.tfs]
        self.avgdl = sum(self.lengths) / max(len(self.lengths), 1)
        # 文档频率：每个词出现在多少段里。
        df = Counter()
        for tf in self.tfs:
            df.update(tf.keys())        # 只加一次：一段里出现多次也只算一段
        n = len(self.passages)
        self.idf = {t: math.log((n - c + 0.5) / (c + 0.5) + 1) for t, c in df.items()}
        # 倒排索引（inverted index）：词 → 含有这个词的段落编号列表（从小到大）。
        # 有了它，检索时只需要给"至少含有一个查询词"的段落打分，不必把几千段全部算一遍。
        # 所有搜索引擎都建立在倒排索引上。（第 2 讲的语料有上千段、要检索几百次，没有它会慢很多。）
        self.postings = {}
        for i, tf in enumerate(self.tfs):
            for t in tf:
                self.postings.setdefault(t, []).append(i)

    def score(self, query_tokens, i):
        """
        查询（已分词）和第 i 段的 BM25 分数。

        【边界】查询里的词如果整个语料都没有，IDF 查不到，对分数没有贡献（贡献为 0）。
        """
        tf, length = self.tfs[i], self.lengths[i]
        s = 0.0
        for t in query_tokens:
            if t not in tf:
                continue
            f = tf[t]
            denom = f + self.k1 * (1 - self.b + self.b * length / self.avgdl)
            s += self.idf[t] * f * (self.k1 + 1) / denom
        return s

    def search(self, query, k=3):
        """
        返回和查询最相关的 k 段。

        【参数】query：str；k：要几段（k ≥ 1）
        【返回】list[Passage]，按分数从高到低。分数全为 0 时（查询里没有一个词出现在语料里）返回空列表。
        【边界】分数相同的段落，按它们在语料里的先后顺序排（Python 的 sorted 是"稳定"排序）。

        【Python 小知识】
            sorted(range(n), key=lambda i: -scores[i])：把编号 0..n−1 按分数从大到小排序。
            lambda i: -scores[i] 是一个没有名字的小函数，输入 i，输出 −scores[i]；
            取负号是因为 sorted 默认从小到大排。
        """
        return [p for p, _score in self.search_scored(query, k)]

    def search_scored(self, query, k=3):
        """
        和 search 一样，但同时返回分数：list[(Passage, 分数)]。

        【思路】先用倒排索引找出"至少含有一个查询词"的候选段落（其他段落的分数一定是 0），
            只给候选打分，再排序取前 k 个。结果和"给全部段落打分"完全一样，只是快得多。
        【边界】含有查询词的段落分数一定大于 0（IDF 和词频都是正数），所以返回的都是分数 > 0 的段落。
        【Python 小知识】
            set().union(*列表的列表)：把几个列表合并成一个集合（去掉重复）；
            sorted(集合) 再把它排成从小到大的列表，保证"分数相同按语料顺序排"。
        """
        q = tokenize(query)
        candidates = sorted(set().union(*[self.postings.get(t, []) for t in q]))
        scores = {i: self.score(q, i) for i in candidates}
        order = sorted(candidates, key=lambda i: -scores[i])
        return [(self.passages[i], scores[i]) for i in order[:k]]


if __name__ == "__main__":
    from toy_world import make_world
    world = make_world(seed=0)
    bm25 = BM25(world.passages)
    d = world.diseases[0].name
    for q in [f"Which gene is mutated in {d}?", f"What is the genetic cause of {d}?"]:
        print(f"\n查询：{q}")
        print(f"  分词：{tokenize(q)}")
        for p in bm25.search(q, k=3):
            print(f"  [{p.pid}] {p.text}")