"""
=============================================================================
第 1 讲 · 代码 01：n-gram 语言模型 —— 数数、稀疏、平滑，以及"看不见的提示"
=============================================================================

对应讲义：lectures/lecture01_history_of_lms.md §2（语言模型是什么：从 Shannon 到 n-gram）

【这个脚本要回答的问题】
    1. "语言模型就是数数"具体是什么意思？给一个提示，n-gram 模型怎样决定下一个词？（演示 1）
    2. 为什么不能把 n 一直加大？看得越远，见过的组合就越少——"稀疏性"。平滑在救什么？（演示 2）
    3. n-gram 模型能"被提示"吗？能，但它只看最后 n−1 个词。窗口外的提示，写了等于没写。（演示 3）
    4. n-gram 模型生成的文字，为什么"每两个词都通顺，整段话却不知所云"？（演示 4）

【怎么读输出】
    每个演示先打印表格，再打印"观察 Observations"。表格里的"困惑度"越低越好，
    直觉是"模型平均在几个词之间犹豫"（这份玩具语料的词表有 290 个词）。

【运行】python3 01_ngram.py      （约 10 秒）
【依赖】只用 Python 标准库；import 同目录的 pretrain_world.py、ngram_lm.py、console.py。
"""

import random

from console import pct, print_header, print_observations, print_table
from ngram_lm import NGramLM, top_k
from pretrain_world import (GENDER_PAIRS, WordTokenizer, build_corpus, build_world, corpus_text,
                            question_text)


def load():
    """
    生成世界、训练语料和测试语料，并编码成 id。

    【返回】(world, tok, train_ids, test_ids, train_docs, test_docs)
    【边界】测试语料用 seed=1：同一个世界（同样的国家、首都），但句子是重新随机生成的。
    """
    world = build_world(seed=0)
    train_docs = build_corpus(world, seed=0)
    test_docs = build_corpus(world, seed=1, scale=0.2)
    text = corpus_text(train_docs)
    tok = WordTokenizer(text)
    return world, tok, tok.encode(text), tok.encode(corpus_text(test_docs)), train_docs, test_docs


def fmt_top(dist, tok, k=3):
    """把分布里前 k 个词写成 "of 0.98, city 0.01" 这样的字符串。"""
    return ", ".join(f"{w} {p:.2f}" for w, p in top_k(dist, tok, k))


# ---------------------------------------------------------------------------
def demo1_counting(world, tok, ids):
    """演示 1：数数——同一个提示，不同的 n 看见的上下文不同，给出的预测也不同。"""
    print_header("演示 1：n-gram 模型怎样决定下一个词——它只看最后 n−1 个词")
    lms = {n: NGramLM(n).fit(ids, len(tok)) for n in [2, 3, 5, 7]}

    # (a) 一个短语的下一词分布：上下文越长，分布越"确定"
    rows = []
    for ctx in ["of", "capital of", "the capital of", "rating :"]:
        n = len(ctx.split()) + 1
        lm = lms.get(n) or NGramLM(n).fit(ids, len(tok))
        rows.append([f"{n}-gram", repr(ctx), fmt_top(lm.dist(tok.encode(ctx)), tok, 4)])
    print_table(["模型", "上下文", "下一个词的前 4 名（概率）"], rows)

    # (b) 问一道知识题：问答格式里见过的国家 vs 只在百科里出现过的国家
    seen, held = world.seen[0], world.held[0]
    print()
    rows = []
    for c in [seen, held]:
        prompt = f"<doc> faq : q : {question_text(c.name, 'capital')} a :"
        for n, lm in lms.items():
            used = tok.decode(lm.context_used(tok.encode(prompt)))
            best = top_k(lm.dist(tok.encode(prompt)), tok, 1)[0]
            ok = "✓" if best[0] == c.capital else "✗"
            rows.append([c.name + ("（问答见过）" if c.qa_seen else "（只在百科）"), f"{n}-gram",
                         repr(used), f"{best[0]} {best[1]:.2f} {ok}"])
    print(f"提示：'<doc> faq : q : what is the capital of <国家> ? a :'    "
          f"标准答案：{seen.name} → {seen.capital}，{held.name} → {held.capital}")
    print_table(["国家", "模型", "模型真正看见的上下文", "预测 top-1"], rows)
    d5 = lms[5].dist(tok.encode(f"<doc> faq : q : {question_text(seen.name, 'capital')} a :"))
    split = {rel: d5[tok.stoi[seen.attr(rel)]] for rel in ["capital", "language", "leader"]}
    print("5-gram 在 " + repr(f"{seen.name} ? a :") + " 之后给三个可能答案的概率："
          + "，".join(f"{rel}={seen.attr(rel)} {p:.2f}" for rel, p in split.items()))
    print_observations([
        "n-gram 的全部'知识'就是一张计数表：P(下一个词 | 前 n−1 个词) = 次数之比。"
        "表 (a) 里 'of' 后面可能是 24 个国家中的任何一个，分布很平——只看这几个词，确实无从知道是哪个国家。",
        f"5-gram 看见的是 '<国家> ? a :'：问首都、问语言、问领导人的问题都以它结尾，所以概率被分给了三种答案"
        f"（首都 {split['capital']:.2f}、语言 {split['language']:.2f}、领导人 {split['leader']:.2f}）；"
        "这次首都最大，碰巧答对。7-gram 同时看见了国家和关系，才给出接近 1 的概率。",
        f"{held.name} 的首都在百科里出现过很多次，但从没以问答格式出现过：7-gram 的上下文在训练里根本不存在，"
        "退回更短的上下文后只剩一个几乎平坦的分布。n-gram 记住的是**字符串**，不是**事实**。",
    ])


# ---------------------------------------------------------------------------
def unseen_rate(lm, test_ids):
    """测试流里，完整的 n 元组（上下文 + 下一个词）在训练中从没出现过的比例。"""
    n = lm.n
    miss = total = 0
    for i in range(n - 1, len(test_ids)):
        h = tuple(test_ids[i - n + 1:i])
        total += 1
        if lm.counts[n - 1].get(h, {}).get(test_ids[i], 0) == 0:
            miss += 1
    return miss / total


def demo2_sparsity(tok, ids, test_ids):
    """演示 2：困惑度随 n 的变化；平滑方法的对比。"""
    print_header("演示 2：看得越远越好吗？——稀疏性与平滑")
    rows = []
    best_n, best_ppl = None, float("inf")
    results = {}
    for n in range(1, 11):
        lm = NGramLM(n, smoothing="kn").fit(ids, len(tok))
        tr, _ = lm.perplexity(ids[:45000])   # 训练语料的前 45000 个 token（和测试集一样长）
        te, _ = lm.perplexity(test_ids)
        un = unseen_rate(lm, test_ids)
        results[n] = (tr, te, un)
        if te < best_ppl:
            best_n, best_ppl = n, te
        rows.append([f"{n}-gram", f"{tr:.3f}", f"{te:.3f}", pct(un, 1)])
    print_table(["模型（Kneser–Ney）", "训练困惑度", "测试困惑度", "测试 n 元组没见过的比例"], rows)

    print()
    rows = []
    for n in [3, 7]:
        for sm, kw in [("mle", {}), ("addk", {"k": 1.0}), ("addk", {"k": 0.01}), ("interp", {}), ("kn", {})]:
            lm = NGramLM(n, smoothing=sm, **kw).fit(ids, len(tok))
            te, zero = lm.perplexity(test_ids)
            name = sm + (f" (k={kw['k']})" if kw else "")
            rows.append([f"{n}-gram", name, "∞" if te == float("inf") else f"{te:.3f}", pct(zero, 1, 2)])
    print_table(["模型", "平滑方法", "测试困惑度", "概率为 0 的测试 token"], rows)
    tr1, te1, _ = results[1]
    print_observations([
        f"1-gram（只看词频）的困惑度约 {te1:.0f}：相当于每一步在 {te1:.0f} 个词里瞎猜。看前 1 个词（2-gram）就骤降到 "
        f"{results[2][1]:.1f}——语言的大部分结构是局部的。",
        f"训练困惑度随 n 一直下降（n=10 时 {results[10][0]:.2f}），测试困惑度在 n={best_n} 时最低（{best_ppl:.3f}），"
        f"之后反而上升：n 越大，测试里没见过的 n 元组越多（n=10 时 {pct(results[10][2], 1)}）。这就是**稀疏性**。",
        "不平滑（mle）时，只要测试里出现一个没见过的组合，概率就是 0、困惑度就是无穷大；n 越大越严重。",
        "加 1 平滑（k=1）把太多概率分给了从没出现过的词，反而很差；插值 / Kneser–Ney 的思路是"
        "'没见过就退回更短的上下文'，效果最好。",
        "这份语料是模板生成的、非常规整，所以 n-gram 在这里表现得比在真实语言上好得多。"
        "真实英语里 5-gram 已经极度稀疏，Google 2007 年用 2 万亿词训练 5-gram 来对付它。",
    ])


# ---------------------------------------------------------------------------
def pronoun_probes(test_docs, tok):
    """
    从测试故事里找出所有 "the <人物> <语境> . <代词>" 的位置：给出前文，问下一个词是 he 还是 she。

    【返回】[(前文 id 列表, 正确代词, 人物词和代词之间隔了几个词), ...]
    【思路】代词由几个词之前的人物决定（queen → she）。隔得越远，越需要长上下文。
        只收"紧跟在 the <人物> ... . 后面"的代词：往前找最近的人物词，要求它前面是 "the"、
        而且它和代词之间只隔着**一个**句号（同一个人物的下一句）。
        bank / vault / shore 那种句子（"she put money in the bank ."）的代词是随机的，不收。
    【边界】第一版直接找"最近的 the"，可 "lives in the palace" 这种语境里本身就有 the，
        把距离算错了；现在按人物词来找。
    """
    persons = {w for pair in GENDER_PAIRS for w in pair}
    probes = []
    for d in test_docs:
        if d.kind != "story":
            continue
        words = ("<doc> " + d.text).split()
        for i, w in enumerate(words):
            if w not in ("he", "she") or words[i - 1] != ".":
                continue
            ks = [k for k in range(i) if words[k] in persons]
            if not ks:
                continue
            k = ks[-1]
            if words[k - 1] == "the" and words[k:i].count(".") == 1:
                probes.append((tok.encode(" ".join(words[:i])), w, i - k))
    return probes


def demo3_window(tok, ids, test_docs):
    """演示 3：提示里窗口之外的部分，n-gram 完全看不见。"""
    print_header("演示 3：n-gram 能被'提示'吗？——窗口外的提示写了等于没写")
    he, she = tok.stoi["he"], tok.stoi["she"]
    probes = pronoun_probes(test_docs, tok)
    rows = []
    for n in [2, 3, 4, 5, 6, 7, 8]:
        lm = NGramLM(n).fit(ids, len(tok))
        ok = sum(1 for ctx, gold, _ in probes
                 if ("he" if lm.prob(he, ctx) > lm.prob(she, ctx) else "she") == gold)
        far = [(ctx, gold) for ctx, gold, dist in probes if dist >= n]   # 人物词在窗口之外
        ok_far = sum(1 for ctx, gold in far if ("he" if lm.prob(he, ctx) > lm.prob(she, ctx) else "she") == gold)
        inside = len(probes) - len(far)
        rows.append([f"{n}-gram", n - 1, pct(inside, len(probes)), pct(ok, len(probes)),
                     f"{pct(ok_far, len(far))}（{len(far)} 个）" if far else "-（0 个）"])
    print(f"任务：故事里 'the queen lives in the palace . ___'，下一个词是 he 还是 she？共 {len(probes)} 个测试位置。")
    print_table(["模型", "窗口（看前几个词）", "人物词在窗口内的比例", "代词准确率", "人物词在窗口之外时的准确率"], rows)

    # 同一个模型，两个只差一个 "not" 的提示：5-gram 给出完全相同的分布
    p_neg = "<doc> reviews : review : the soup was not cold . rating :"
    p_pos = "<doc> reviews : review : the soup was cold . rating :"
    print()
    rows = []
    for n in [5, 7]:
        lm = NGramLM(n).fit(ids, len(tok))
        d1, d2 = lm.dist(tok.encode(p_neg)), lm.dist(tok.encode(p_pos))
        diff = max(abs(a - b) for a, b in zip(d1, d2))
        rows.append([f"{n}-gram", repr(tok.decode(lm.context_used(tok.encode(p_neg)))),
                     fmt_top(d1, tok, 2), fmt_top(d2, tok, 2), f"{diff:.3f}"])
    print(f"提示 A：{p_neg!r}（正确：positive）\n提示 B：{p_pos!r}（正确：negative）")
    print_table(["模型", "A 的可见上下文", "A 的预测", "B 的预测", "两个分布的最大差"], rows)
    print_observations([
        "人物词在代词之前 4–6 个位置（'the queen lives in the palace . she'，中间隔着 3–5 个词）。人物词落在窗口之外时，"
        "n-gram 只能按 he / she 的整体频率去猜，准确率在一半左右；窗口一旦盖住人物词，准确率就跳到接近 100%。",
        "5-gram 对'not cold'和'cold'给出的分布**一模一样**（最大差 0）：'not' 在它的窗口外面。7-gram 才看得见。",
        "这就是为什么在 n-gram 时代没有'提示工程'：你在提示开头写'请判断下面评论的情感'，模型一个字都看不到。"
        "能利用整段提示的长距离依赖，是神经网络语言模型——尤其是 Transformer——带来的。",
    ])


# ---------------------------------------------------------------------------
def story_metrics(gens, tok, story_vocab):
    """
    给一批生成的"故事"打两个分。

    【参数】gens：生成的 id 列表的列表；story_vocab：训练语料的故事里出现过的所有词
    【返回】(没跑题的比例, 代词和人物性别一致的比例, 可检查的代词个数)
    【思路】
        没跑题：生成的每个词都在故事词表里（一旦冒出 "?"、"rating"、"capital"，说明它"写着写着
            变成了别的文档"——局部通顺，整体不连贯）。
        代词一致：找 "the <人物> ... . <代词>"，看 queen 后面是不是 she、king 后面是不是 he。
    """
    female = {f for _, f in GENDER_PAIRS}
    male = {m for m, _ in GENDER_PAIRS}
    on_topic = agree = checked = 0
    for g in gens:
        words = tok.decode(g).split()
        on_topic += all(w in story_vocab for w in words)
        for i, w in enumerate(words):
            if w in ("he", "she") and i >= 1 and words[i - 1] == ".":
                ks = [k for k in range(i) if words[k] in male | female]
                if ks and words[ks[-1] - 1] == "the" and words[ks[-1]:i].count(".") == 1:
                    checked += 1
                    agree += (w == "she") == (words[ks[-1]] in female)
    return on_topic / len(gens), (agree / checked if checked else 0.0), checked


def demo4_generation(tok, ids, train_docs):
    """演示 4：用 n-gram 模型生成文字——局部通顺，不等于整体连贯。"""
    print_header("演示 4：n-gram 生成的文字——局部通顺，不等于整体连贯")
    story_vocab = {w for d in train_docs if d.kind == "story" for w in d.text.split()}
    doc = tok.stoi["<doc>"]
    rows, samples, agree_by_n = [], {}, {}
    for n in [2, 3, 5, 8]:
        lm = NGramLM(n).fit(ids, len(tok))
        rng = random.Random(n)
        gens = [lm.generate(tok.encode("<doc> story :"), 40, temperature=1.0, rng=rng, stop=doc) for _ in range(200)]
        samples[n] = tok.decode(gens[0])
        topic, agree, checked = story_metrics(gens, tok, story_vocab)
        agree_by_n[n] = agree
        rows.append([f"{n}-gram", pct(topic, 1), f"{pct(agree, 1)}（{checked} 处）"])
    print_table(["模型（温度 1.0，各生成 200 段）", "整段没跑题的比例", "代词与人物性别一致"], rows)
    print()
    for n, s in samples.items():
        print(f"  {n}-gram：story : {s}")
    print_observations([
        "2-gram 的每两个相邻的词都'说得通'，但整段话很快就跑题：一个故事写着写着变成了题目列表、评论……"
        "它只记得上一个词，不记得自己在写哪种文档。",
        f"代词一致需要记住 4–6 个位置之前的人物词：2-gram、3-gram、5-gram 的一致率只有 {pct(agree_by_n[2], 1)} / "
        f"{pct(agree_by_n[3], 1)} / {pct(agree_by_n[5], 1)}（≈ 瞎猜），窗口盖住人物词的 8-gram 是 {pct(agree_by_n[8], 1)}。",
        "n 越大，生成的文字越像训练语料——但代价是演示 2 里的稀疏性。n-gram 被困在这个两难里："
        "看得近就不连贯，看得远就没见过。神经网络语言模型（§5）的出路是：不再把上下文当成一个必须原样见过的字符串。",
    ])


def main():
    world, tok, ids, test_ids, train_docs, test_docs = load()
    print(f"训练语料 {len(ids):,} 个 token，测试语料 {len(test_ids):,} 个 token，词表 {len(tok)} 个词")
    demo1_counting(world, tok, ids)
    demo2_sparsity(tok, ids, test_ids)
    demo3_window(tok, ids, test_docs)
    demo4_generation(tok, ids, train_docs)


if __name__ == "__main__":
    main()