"""
=============================================================================
CMU 17-630 · 第 1 讲共享模块：ngram_lm.py —— 数数就能做的语言模型
=============================================================================

【这个模块是做什么的】
    语言模型（language model, LM）就是一个给"下一个词"打概率的函数：

        P(w_t | w_1, ..., w_{t-1})

    从 1940 年代 Shannon 的实验到 2010 年代初，最成功的语言模型都是 **n-gram 模型**：
    假设下一个词只取决于前面 n−1 个词（"马尔可夫假设"），然后**数数**：

        P(w | 前 n−1 个词 h) ≈ count(h, w) / count(h)

    这个模块实现了它，以及三种"平滑"（smoothing）方法。平滑要解决的问题是：
    测试时总会遇到训练里没出现过的 n 元组，如果直接数数，概率就是 0，
    一句话里只要有一个 0，整句话的概率就是 0，困惑度就是无穷大。

        mle     最大似然：直接数数，不平滑（会出现 0 概率）
        addk    加 k 平滑：每个计数都加一个小数 k（简单，但效果一般）
        interp  插值的绝对折扣（absolute discounting）：从每个见过的计数里扣掉一点 d，
                把扣下来的概率按"更短上下文"的模型分出去
        kn      Kneser–Ney：在 interp 的基础上，低阶模型用"这个词跟在多少种不同的词后面"
                （continuation count）代替原始计数。n-gram 时代公认最好的平滑方法之一。

【为什么提示工程的课要讲它】
    n-gram 模型能"被提示"：给它一段前文，它就续写。但它**只看最后 n−1 个词**——
    你在提示开头写的任何指令，它都完全看不见（01 号脚本演示 3）。
    "长距离地利用整个提示"正是后来的神经网络语言模型、尤其是 Transformer 带来的能力，
    也是"提示工程"这件事能存在的前提。

【运行】python3 ngram_lm.py     # 自检：在训练语料上拟合，打印几个下一词分布
【依赖】只用 Python 标准库。
"""

import math
import random
from collections import Counter, defaultdict


class NGramLM:
    """
    n-gram 语言模型（作用在 token id 上）。

    【用法】
        lm = NGramLM(n=3, smoothing="kn").fit(ids, vocab_size)
        lm.prob(w, context)        # 单个词的条件概率
        lm.dist(context)           # 整个词表上的分布（列表）
        lm.perplexity(test_ids)    # 困惑度
        lm.generate(context, 20)   # 续写
    【参数】
        n：阶数（n=1 只看词频，n=2 看前 1 个词，n=3 看前 2 个词……）
        smoothing："mle" / "addk" / "interp" / "kn"
        k：addk 用的加数；discount：interp / kn 用的折扣 d（0 < d < 1，经验上 0.75 左右最好）
    """

    def __init__(self, n, smoothing="kn", k=0.1, discount=0.75):
        assert n >= 1 and smoothing in ("mle", "addk", "interp", "kn")
        self.n, self.smoothing, self.k, self.d = n, smoothing, k, discount

    # ------------------------------------------------------------------ 训练：数数
    def fit(self, ids, vocab_size):
        """
        在一条 token id 流上数出所有 1…n 元组的次数。

        【参数】ids：整数列表（整份训练语料）；vocab_size：词表大小 V
        【返回】self（这样可以写成 NGramLM(3).fit(...) 一行）
        【思路】counts[m][h][w] = "上下文 h（长度 m）后面紧跟着 w" 出现的次数，m = 0 … n−1。
            m = 0 时 h 是空元组 ()，counts[0][()] 就是词频。
            Kneser–Ney 还需要 cont[m][h][w] = "有多少种不同的词 u，使得 (u, *h, w) 出现过"。
        【Python 小知识】
            defaultdict(Counter)：访问一个不存在的键时自动创建一个空 Counter，省去 "if key not in"。
            tuple(ids[i-m:i])：列表切片再转成元组——元组不可变，所以能当字典的键。
        """
        self.V = vocab_size
        self.counts = [defaultdict(Counter) for _ in range(self.n)]
        for i in range(len(ids)):
            w = ids[i]
            for m in range(self.n):
                if i - m < 0:
                    break
                self.counts[m][tuple(ids[i - m:i])][w] += 1
        # 预先算好每个上下文的总次数和"后面跟过几种不同的词"，避免每次查概率都重新求和
        self.totals = [{h: sum(c.values()) for h, c in level.items()} for level in self.counts]
        self.types = [{h: len(c) for h, c in level.items()} for level in self.counts]
        if self.smoothing == "kn":
            # cont[m][h][w]：(m+1)+1 元组 (u, *h, w) 里有多少种不同的 u
            self.cont = [defaultdict(Counter) for _ in range(self.n)]
            for m in range(self.n - 1):
                for h_long, c in self.counts[m + 1].items():   # h_long 长度 m+1 = (u, *h)
                    for w in c:
                        self.cont[m][h_long[1:]][w] += 1
            self.cont_totals = [{h: sum(c.values()) for h, c in lv.items()} for lv in self.cont]
            self.cont_types = [{h: len(c) for h, c in lv.items()} for lv in self.cont]
        return self

    # ------------------------------------------------------------------ 查概率
    def _level_prob(self, w, h, top):
        """
        在上下文 h（元组，长度 m）这一层计算 P(w | h)，必要时递归到更短的上下文 h[1:]。

        【参数】w：词 id；h：上下文元组；top：是不是最外层（Kneser–Ney 只在最外层用原始计数）
        【返回】概率（0–1 之间的小数）
        【思路】以 interp 为例：
            P(w | h) = max(c(h,w) − d, 0) / c(h)  +  d · T(h) / c(h) · P(w | h[1:])
            第一项：见过的就用计数，但每个计数扣掉 d；
            第二项：扣下来的总量正好是 d · T(h)（T(h) 是 h 后面出现过几种不同的词），
                    按更短上下文的分布分给所有词——包括没见过的词，于是不会有 0。
            如果 h 根本没出现过（c(h) = 0），直接退到更短的上下文。
            最底层（m = 0）退到均匀分布 1/V。
        """
        m = len(h)
        use_cont = self.smoothing == "kn" and not top
        tables = (self.cont, self.cont_totals, self.cont_types) if use_cont else (self.counts, self.totals, self.types)
        counts, totals, types = tables[0][m], tables[1][m], tables[2][m]
        c_h = totals.get(h, 0)
        if self.smoothing == "mle":
            return counts[h][w] / c_h if c_h else 0.0
        if self.smoothing == "addk":
            return (counts[h][w] + self.k) / (c_h + self.k * self.V) if c_h else 1.0 / self.V
        # interp / kn
        lower = self._level_prob(w, h[1:], False) if m > 0 else 1.0 / self.V
        if c_h == 0:
            return lower
        c_hw = counts[h].get(w, 0) if h in counts else 0
        return max(c_hw - self.d, 0.0) / c_h + self.d * types[h] / c_h * lower

    def context_used(self, context):
        """
        模型真正"看得见"的那部分上下文：最后 n−1 个 token（不足时就是全部）。

        【提示工程联系】这就是 n-gram 模型的"上下文窗口"。窗口之外的提示内容对输出**没有任何影响**。
        """
        context = list(context)
        return tuple(context[-(self.n - 1):]) if self.n > 1 else ()

    def prob(self, w, context):
        """P(w | context)。context 可以任意长，只有最后 n−1 个 token 会被用到。"""
        return self._level_prob(w, self.context_used(context), True)

    def dist(self, context):
        """
        整个词表上的下一词分布。

        【返回】长度为 V 的概率列表（mle 在没见过的上下文上会全是 0）
        【边界】interp / kn / addk 的结果加起来等于 1（浮点误差内）；mle 在见过的上下文上也等于 1。
        """
        h = self.context_used(context)
        return [self._level_prob(w, h, True) for w in range(self.V)]

    # ------------------------------------------------------------------ 评估：困惑度
    def perplexity(self, ids):
        """
        在一条测试 token 流上的困惑度（perplexity）。

        【思路】困惑度 = exp(平均每个 token 的负对数概率)。直觉："模型平均在多少个词之间犹豫"。
            一个完全随机的模型（均匀分布）困惑度 = V；完美预测的模型 = 1。
        【返回】(困惑度, 概率为 0 的 token 比例)
        【边界】出现概率为 0 的 token 时（只有 mle 会这样），困惑度是无穷大 math.inf；
            第一个 token 没有上下文，跳过。
        """
        nll, zeros, n = 0.0, 0, 0
        for i in range(1, len(ids)):
            p = self.prob(ids[i], ids[max(0, i - self.n + 1):i])
            n += 1
            if p <= 0.0:
                zeros += 1
            else:
                nll -= math.log(p)
        if zeros:
            return math.inf, zeros / n
        return math.exp(nll / n), 0.0

    # ------------------------------------------------------------------ 生成：采样续写
    def generate(self, context, n_tokens, temperature=1.0, rng=None, stop=None):
        """
        从 context 开始续写 n_tokens 个 token。

        【参数】context：起始 token id 列表；temperature：温度（见下）；rng：random.Random；
            stop：遇到这个 id 就停（例如 "<doc>"）
        【返回】新生成的 id 列表（不含 context）
        【思路】温度 T：把每个概率变成 p^(1/T) 再归一化。T < 1 让分布更"尖"（更保守），
            T > 1 让分布更"平"（更随机），T → 0 等价于每次取概率最大的词（贪心，greedy）。
        【边界】temperature <= 0 时直接取 argmax；mle 遇到没见过的上下文（全 0 分布）时停止。
        """
        rng = rng or random.Random(0)
        ctx = list(context)
        out = []
        for _ in range(n_tokens):
            p = self.dist(ctx)
            if sum(p) <= 0:
                break
            if temperature <= 0:
                w = max(range(self.V), key=lambda i: p[i])
            else:
                weights = [x ** (1.0 / temperature) for x in p]
                w = rng.choices(range(self.V), weights=weights)[0]
            if stop is not None and w == stop:
                break
            out.append(w)
            ctx.append(w)
        return out


def top_k(dist, tok, k=5):
    """
    把一个分布里概率最大的 k 个词挑出来，便于打印。

    【参数】dist：概率列表；tok：WordTokenizer；k：个数
    【返回】[(词, 概率), ...]，按概率从大到小
    【Python 小知识】sorted(..., key=lambda i: -dist[i]) 按概率从大到小排序；
        lambda 是"一次性的小函数"，lambda i: -dist[i] 等价于 def f(i): return -dist[i]。
    """
    order = sorted(range(len(dist)), key=lambda i: (-dist[i], i))[:k]
    return [(tok.itos[i], dist[i]) for i in order]


if __name__ == "__main__":
    from pretrain_world import WordTokenizer, build_corpus, build_world, corpus_text

    world = build_world()
    text = corpus_text(build_corpus(world, seed=0))
    tok = WordTokenizer(text)
    ids = tok.encode(text)
    test_ids = tok.encode(corpus_text(build_corpus(world, seed=1, scale=0.2)))
    for sm in ["mle", "addk", "interp", "kn"]:
        lm = NGramLM(3, smoothing=sm).fit(ids, len(tok))
        ppl, z = lm.perplexity(test_ids)
        s = sum(lm.dist(tok.encode("the capital of")))
        print(f"3-gram {sm:6s}  测试困惑度 {ppl:8.2f}  零概率 token {100 * z:5.2f}%  分布之和 {s:.6f}")
    lm = NGramLM(3).fit(ids, len(tok))
    for prompt in ["the capital", "review : the", "rating :"]:
        tops = ", ".join(f"{w} {p:.2f}" for w, p in top_k(lm.dist(tok.encode(prompt)), tok))
        print(f"P(下一个词 | {prompt!r}) 前 5：{tops}")