"""
=============================================================================
第 1 讲 · 代码 1：语言模型是什么？—— 序列上的概率分布、温度、n-gram 与数据稀疏
=============================================================================

对应讲义：lectures/lecture01_intro_foundation_models.md 第 2 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
"基础模型"的核心是语言模型（language model, LM）。在讨论 GPT-3 有多神奇之前，先把最朴素的定义弄清楚：

    演示 1：语言模型 = 序列上的概率分布。用 CS324 课程笔记里的 5 个词（ate, ball, cheese, mouse, the）
            手写一个小模型：算 p(the mouse ate the cheese)，用链式法则逐项分解；
            验证"所有可能的句子"的概率加起来等于 1（结束符 </s> 是关键）；
            贪心生成为什么会陷入重复。
    演示 2：温度（temperature）。复现课程笔记里 0.4 / 0.6 的例子；
            再用一个两步的小例子说明："每一步都调温度"和"对整句的分布调温度"不是一回事，
            所以贪心解码（每步选最可能的词）不一定得到最可能的句子。
    演示 3：n-gram 语言模型 + 数据稀疏。在迷你互联网（fm_world.py）上训练 1–10 阶字符 n-gram：
            n 越大，训练集上越好，验证集上先好后坏；没见过的上下文越来越多。
    演示 4：从 n-gram 里采样。低阶：乱码；高阶：格式完美，但很多句子是训练集里的原句（背诵），
            而"新"句子常常是假的（幻觉）。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 01_what_is_a_language_model.py          # 约 3 秒

只用 Python 标准库。依赖本目录的 fm_world.py、ngram_lm.py、console.py。
"""

import itertools
import math
import random

from console import pad, print_header, show
from fm_world import ALPHABET, build_world, check_line
from ngram_lm import CharNGramLM

# =============================================================================
# 演示 1：语言模型 = 序列上的概率分布
# =============================================================================
END = "</s>"   # 结束符：句子到这里结束
WORDS = ["ate", "ball", "cheese", "mouse", "the"]

# 一个手写的 bigram 模型：第一个词的分布 + "上一个词 → 下一个词"的分布。
# 每一行的概率加起来都是 1。数字是我随手设定的，只要"像英语"即可：
#   "the" 后面多半是名词；"mouse" 后面多半是 "ate"；"cheese" / "ball" 后面多半结束；"ate" 后面多半是 "the"。
FIRST = {"the": 0.70, "mouse": 0.10, "cheese": 0.08, "ball": 0.07, "ate": 0.05}
NEXT = {
    "the":    {"mouse": 0.40, "cheese": 0.35, "ball": 0.20, "ate": 0.04, "the": 0.01, END: 0.00},
    "mouse":  {"ate": 0.60, END: 0.25, "the": 0.05, "ball": 0.04, "cheese": 0.03, "mouse": 0.03},
    "cheese": {END: 0.70, "the": 0.10, "ate": 0.05, "ball": 0.05, "mouse": 0.05, "cheese": 0.05},
    "ball":   {END: 0.70, "the": 0.10, "ate": 0.05, "mouse": 0.05, "cheese": 0.05, "ball": 0.05},
    "ate":    {"the": 0.80, END: 0.10, "cheese": 0.05, "mouse": 0.02, "ball": 0.02, "ate": 0.01},
}


def sentence_prob(words, verbose=False):
    """
    用链式法则算一个句子的概率：p(x1..xL) = p(x1) · p(x2|x1) · ... · p(xL|x_{L−1}) · p(</s>|xL)。

    【思路】这是一个 bigram 模型，所以 p(x_i | x_1..x_{i−1}) 只取决于上一个词 x_{i−1}。
            最后乘上"在这里结束"的概率 p(</s> | x_L)——没有这一项，"the" 和 "the mouse" 的概率
            就没法放进同一个分布里比较（见 demo_lm_is_a_distribution 的第 (b) 部分）。
    【参数】words：词的列表；verbose：True 时打印每一项
    【返回】概率（浮点数）
    【边界】空列表返回 0（我们的模型不生成空句子）；出现 WORDS 之外的词时 KeyError。
    """
    if not words:
        return 0.0
    p = FIRST[words[0]]
    terms = [f"p({words[0]}) = {FIRST[words[0]]:.2f}"]
    for prev, cur in zip(words, words[1:]):
        p *= NEXT[prev][cur]
        terms.append(f"p({cur}|{prev}) = {NEXT[prev][cur]:.2f}")
    p *= NEXT[words[-1]][END]
    terms.append(f"p(</s>|{words[-1]}) = {NEXT[words[-1]][END]:.2f}")
    if verbose:
        print("    = " + " × ".join(terms))
    return p


def sample_words(rng, prompt, temperature, max_words=12):
    """
    从手写的 bigram 模型里续写 prompt。

    【参数】rng：random.Random；prompt：已有的词列表；temperature：0 = 贪心；max_words：最长生成几个词
    【返回】生成的词列表（不含 prompt；如果生成了 </s> 就在那里停）
    """
    out = []
    prev = prompt[-1]
    for _ in range(max_words):
        options = list(NEXT[prev].items())
        if temperature == 0:
            w = max(options, key=lambda kv: kv[1])[0]
        else:
            weights = [p ** (1 / temperature) for _, p in options]
            w = rng.choices([w for w, _ in options], weights=weights)[0]
        if w == END:
            break
        out.append(w)
        prev = w
    return out


def demo_lm_is_a_distribution():
    print_header("演示 1：语言模型 = 序列上的概率分布（词表 = ate, ball, cheese, mouse, the）")

    print("(a) 三个句子的概率（链式法则逐项相乘）：")
    for s in ["the mouse ate the cheese", "the cheese ate the mouse", "mouse the the cheese ate"]:
        p = sentence_prob(s.split(), verbose=False)
        print(f"  p({s}) = {p:.3g}")
        sentence_prob(s.split(), verbose=True)

    print("\n(b) 所有句子的概率加起来是 1 吗？长度为 L 的句子有 5^L 个，逐个枚举再求和：")
    # 动态规划：alive[w] = "已经生成了 L 个词、最后一个是 w、还没结束"的概率
    alive = dict(FIRST)
    cumulative = 0.0
    print("  " + pad("长度 L", 8) + pad("句子个数", 10, ">") + pad("长度恰为 L 的总概率", 22, ">")
          + pad("累计（长度 ≤ L）", 20, ">") + pad("枚举验证", 14, ">"))
    for L in range(1, 11):
        exact_L = sum(alive[w] * NEXT[w][END] for w in WORDS)
        cumulative += exact_L
        check = ""
        if L <= 5:
            # 【Python 小知识】itertools.product(WORDS, repeat=L) 生成所有长度为 L 的词组合（5^L 个）
            brute = sum(sentence_prob(list(ws)) for ws in itertools.product(WORDS, repeat=L))
            check = "一致 ✓" if abs(brute - exact_L) < 1e-12 else f"不一致 {brute}"
        print("  " + pad(L, 8) + pad(f"{5 ** L:,}", 10, ">") + pad(f"{exact_L:.4f}", 22, ">")
              + pad(f"{cumulative:.4f}", 20, ">") + pad(check, 14, ">"))
        alive = {w2: sum(alive[w] * NEXT[w][w2] for w in WORDS) for w2 in WORDS}

    print("\n(c) 条件生成：提示 prompt = 'the mouse ate'")
    rng = random.Random(0)
    print("  温度 0（贪心）：the mouse ate | " + " ".join(sample_words(rng, ["the", "mouse", "ate"], 0)) + " ……")
    for i in range(4):
        print(f"  温度 1，第 {i + 1} 次：the mouse ate | " + " ".join(sample_words(rng, ["the", "mouse", "ate"], 1.0)))
    print("""
观察 Observations：
  · 同一个模型给"the mouse ate the cheese"的概率是"the cheese ate the mouse"的 34 倍，
    是乱序句子"mouse the the cheese ate"的几十万倍。要把概率分配好，模型得"懂"一点语法（"the the"几乎不可能）
    和一点世界知识（老鼠吃奶酪，而不是奶酪吃老鼠）。这就是为什么"把语言模型做好"会顺带学到很多东西。
  · 长度恰为 L 的句子概率越来越小，但累计概率一步步逼近 1：语言模型是**所有有限长度句子**上的一个分布。
    结束符 </s> 是关键——没有它，长度 1 的句子已经占满了全部概率，"the mouse"就无处安放。
  · 贪心解码（每一步都选概率最大的词）陷入了循环："the → mouse → ate → the → ……"，永远不结束。
    每一步都"最合理"，合起来却是一个很不像样的句子。真实的 LLM 用贪心解码也会出现重复，
    这是采样（温度 > 0）、top-p、重复惩罚等解码技巧存在的原因之一。""")


# =============================================================================
# 演示 2：温度
# =============================================================================
def anneal(probs, T):
    """
    对一个分布做温度调整：p_T(x) ∝ p(x)^(1/T)。

    【参数】probs：dict 或 list 形式的概率；T：温度（> 0），T = float("inf") 表示均匀分布
    【返回】与输入同类型的新分布（已归一化）
    【边界】T = 0 没有定义（1/0），用 anneal_limit_zero() 处理。
    """
    items = list(probs.items()) if isinstance(probs, dict) else list(enumerate(probs))
    weights = [1.0 if math.isinf(T) else p ** (1.0 / T) for _, p in items]
    Z = sum(weights)
    out = [(k, w / Z) for (k, _), w in zip(items, weights)]
    return dict(out) if isinstance(probs, dict) else [v for _, v in out]


def anneal_limit_zero(probs):
    """T → 0 的极限：所有概率集中到最大的那一项（若有并列，平分）。"""
    m = max(probs.values())
    winners = [k for k, v in probs.items() if v == m]
    return {k: (1.0 / len(winners) if k in winners else 0.0) for k in probs}


def demo_temperature():
    print_header("演示 2：温度 Temperature —— p_T(x) ∝ p(x)^(1/T)")
    base = {"cheese": 0.4, "mouse": 0.6}
    print("(a) 课程笔记里的例子：p(cheese) = 0.4，p(mouse) = 0.6")
    for T in [float("inf"), 2.0, 1.0, 0.5, 0.2]:
        p = anneal(base, T)
        print(f"  T = {T:>4}：p(cheese) = {p['cheese']:.2f}，p(mouse) = {p['mouse']:.2f}")
    p = anneal_limit_zero(base)
    print(f"  T → 0 ：p(cheese) = {p['cheese']:.2f}，p(mouse) = {p['mouse']:.2f}   （贪心）")

    print("\n(b) 两步的小模型：先生成 A / B，再生成 x / y")
    first = {"A": 0.6, "B": 0.4}
    second = {"A": {"x": 0.5, "y": 0.5}, "B": {"x": 0.9, "y": 0.1}}
    joint = {a + b: first[a] * second[a][b] for a in first for b in "xy"}
    print("  原始的整句分布：" + "，".join(f"p({k}) = {v:.2f}" for k, v in joint.items())
          + "   ← 最可能的整句是 Bx")
    for T in [0.5, 0.1]:
        joint_T = anneal(joint, T)
        first_T = anneal(first, T)
        step_T = {a + b: first_T[a] * anneal(second[a], T)[b] for a in first for b in "xy"}
        print(f"  T = {T}：对整句调温度  " + "，".join(f"{k} {v:.3f}" for k, v in joint_T.items()))
        print(f"          每步分别调温度 " + "，".join(f"{k} {v:.3f}" for k, v in step_T.items()))
    print("""
观察 Observations：
  · T < 1 让分布更"尖"（高概率的更高），T > 1 让分布更"平"，T → ∞ 是均匀分布，T → 0 是贪心。
    温度只改"形状"，不改排序：概率最大的那个永远还是最大的。
  · 实际生成时，我们**每一步**对下一个 token 的分布调温度——这和"对整句的分布调温度"不一样（只有 T = 1 时相同）。
    在这个例子里，整句最可能的是 Bx（0.36），但每步调温度时 A 在第一步就占了上风：
    T = 0.1 时整句分布几乎全给了 Bx，逐步调温度却几乎全给了 Ax / Ay。
  · 推论：**贪心解码不等于"找最可能的句子"**。第一步选了局部最优的 A，就再也到不了全局最优的 Bx。
    beam search 等方法就是为了缓解这一点；但对开放式生成，"最可能的句子"本身也往往不是最好的句子。""")


# =============================================================================
# 演示 3：n-gram 语言模型与数据稀疏
# =============================================================================
def genre_unseen_rate(world, lm, genre):
    """
    验证集里属于某个体裁的那些行，有多大比例的字符位置的上下文在训练语料里没出现过。

    【参数】world：World；lm：训练好的 CharNGramLM；genre：体裁名
    【返回】0 到 1 之间的比例
    """
    text = "".join(line + "\n" for line in world.val_lines if world.line_genres[line] == genre)
    return lm.unseen_context_rate(text)


def demo_ngram_sparsity(world):
    print_header("演示 3：字符级 n-gram 语言模型 —— n 越大越好吗？")
    train, val = world.train_text(), world.val_text()
    print(f"训练语料 {len(train):,} 个字符，验证语料 {len(val):,} 个字符，字母表 |V| = {len(ALPHABET)}")
    print(f"参照：均匀分布（什么都没学）= log2({len(ALPHABET)}) = {math.log2(len(ALPHABET)):.3f} bit/字符\n")
    print(pad("n", 4) + pad("不同上下文个数", 16, ">") + pad("训练 bit/字符", 15, ">") + pad("验证 bit/字符", 15, ">")
          + pad("验证困惑度", 12, ">") + pad("没见过的上下文：全部", 22, ">") + pad("其中 math 行", 14, ">")
          + pad("其余体裁", 10, ">"))
    results = {}
    for n in [1, 2, 3, 4, 5, 6, 8, 10, 14, 20, 24, 32]:
        lm = CharNGramLM(n, ALPHABET, k=0.01).fit(train)
        tr, va = lm.bits_per_char(train), lm.bits_per_char(val)
        results[n] = va
        other = max(genre_unseen_rate(world, lm, g) for g in ["bio", "qa", "review", "story"])
        print(pad(n, 4) + pad(f"{len(lm.totals):,}", 16, ">") + pad(f"{tr:.3f}", 15, ">")
              + pad(f"{va:.3f}", 15, ">") + pad(f"{2 ** va:.2f}", 12, ">")
              + pad(f"{lm.unseen_context_rate(val):.1%}", 22, ">")
              + pad(f"{genre_unseen_rate(world, lm, 'math'):.1%}", 14, ">") + pad(f"≤ {other:.1%}", 10, ">"))
    best = min(results, key=results.get)
    print(f"\n验证集上最好的是 n = {best}（{results[best]:.3f} bit/字符）。")
    print("""
观察 Observations：
  · 困惑度（perplexity）= 2^(bit/字符)，可以理解为"模型平均在几个字符之间犹豫"。
    均匀分布是 44 个；unigram（只知道字母频率）约 23 个；看前 3 个字符（n = 4）时已经不到 2 个。
  · 训练集上，n 越大越好（看得越远，猜得越准，直到开始"背整行"）；验证集上，n = 14 左右最好，
    再往上（n = 24、32）反而变差：上下文越长，越多的上下文在训练语料里**一次都没出现过**，
    模型在那里只能给出均匀分布（每个字符 5.46 bit）。这就是 n-gram 的根本矛盾——
    n 小，看不到远处的信息；n 大，数据不够（数据稀疏，data sparsity）。
  · 稀疏集中在哪里？看最后两列：**几乎全部出现在 math 行**。人物介绍、评论、故事是由很少的词按模板拼出来的，
    10 个字符的上下文早就在训练集里见过；而两位数加法有 8100 道题，训练集只抽了约 1900 行，
    一道新题的"37+45="在 n = 10 时就是没见过的上下文——math 行里超过 40% 的位置都是这样。
  · 真实的语言更像我们的 math 体裁，而不是模板体裁：词的组合几乎是无穷的。所以在真实英语上，
    稀疏在 n = 4、5 时就已经很严重，即使语料有万亿个词。神经语言模型的出路是：不再"按字面查表"，
    而是把相似的上下文映射到相近的向量上，让"没见过但相似"的上下文也能共享统计信息
    （第 1 讲 §2.5；代码 02 的 base model 在同一份数据上做到了多少，见那里的表格）。""")
    return results


# =============================================================================
# 演示 4：从 n-gram 里采样 —— 乱码、背诵与幻觉
# =============================================================================
def sample_quality(world, lm, rng, n_samples=300, temperature=1.0):
    """
    从模型里采样 n_samples 行，统计格式合法率、为真率、"一字不差地照抄训练集"的比例。

    【思路】每次从 "\\n"（新文档的开头）开始生成，到下一个 "\\n" 为止，得到一行；
            用 fm_world.check_line 判断格式和真假；用集合判断是否是训练集原句。
    【返回】dict：well_formed, true_given_formed, copied, novel_formed_true, examples
            - well_formed       ：格式合法的行占全部行的比例
            - true_given_formed ：格式合法的行里，内容为真的比例
            - copied            ：和训练集某一行一字不差的比例
            - novel_true        ：既不是照抄、又格式合法且为真的比例（真正"生成"出的好句子）
            - fact_true, n_facts：格式合法的人物介绍 / 问答行里，事实为真的比例，以及这样的行有几行
    【边界】100 个字符之内没生成出 "\\n" 的，算作一行格式不合法的文本。
    """
    train_set = set(world.train_lines)
    lines = []
    for _ in range(n_samples):
        s = lm.sample(rng, prefix="\n", n_chars=100, temperature=temperature)
        lines.append(s[:-1] if s.endswith("\n") else s + "…")
    checks = [check_line(world, line) for line in lines]
    n_formed = sum(1 for _, f, _ in checks if f)
    n_true = sum(1 for _, f, t in checks if f and t)
    copied = sum(1 for line in lines if line in train_set)
    novel_true = sum(1 for line, (_, f, t) in zip(lines, checks) if f and t and line not in train_set)
    facts = [t for g, f, t in checks if f and g in ("bio", "qa")]
    return dict(well_formed=n_formed / n_samples, true_given_formed=n_true / max(1, n_formed),
                copied=copied / n_samples, novel_true=novel_true / n_samples,
                fact_true=sum(facts) / max(1, len(facts)), n_facts=len(facts), examples=lines[:3])


def demo_ngram_samples(world):
    print_header("演示 4：从 n-gram 模型里采样（温度 1，每个模型 300 行）")
    train = world.train_text()
    print(pad("n", 4) + pad("格式合法", 10, ">") + pad("其中内容为真", 14, ">") + pad("照抄训练集原句", 16, ">")
          + pad("新的、合法且为真", 18, ">") + pad("人物事实为真", 16, ">"))
    samples = {}
    for n in [1, 2, 3, 4, 6, 8, 10, 14]:
        lm = CharNGramLM(n, ALPHABET, k=0.01).fit(train)
        q = sample_quality(world, lm, random.Random(100 + n))
        samples[n] = q["examples"]
        print(pad(n, 4) + pad(f"{q['well_formed']:.0%}", 10, ">") + pad(f"{q['true_given_formed']:.0%}", 14, ">")
              + pad(f"{q['copied']:.0%}", 16, ">") + pad(f"{q['novel_true']:.0%}", 18, ">")
              + pad(f"{q['fact_true']:.0%}（{q['n_facts']} 行）", 16, ">"))
    print("\n样例（每个模型的前 3 行；…… 表示 100 个字符内没写完一行）：")
    for n in [2, 4, 10]:
        for e in samples[n]:
            print(f"  n = {n:<3} {show(e)[:90]}")
    print("""
观察 Observations：
  · 低阶（n ≤ 3）生成的基本是乱码：局部的字母组合像样，整行几乎没有合法的句子。
  · n = 4 时有约四分之一的行"格式合法"，但几乎全是**算错的加法**（例如 53+95=120.）：
    "两个数、加号、等号、句号"这种格式只靠前 3 个字符就能学会，算对和却需要看到整道题。
  · n 再大，照抄训练集原句的比例一路上升（n = 14 时三分之一）：模型在**背诵**。
    这里的"背诵"在第 16 讲（隐私：记忆与训练数据提取）里会变成一个严肃的问题。
  · 一半左右的行在中途"掉下悬崖"（见 n = 10 的第 3 个样例：60+82=142. 后面全是乱码）：
    生成到一个训练集里没见过的上下文（最常见的是一道新的加法题），模型只能给均匀分布，
    抽出一个怪字符之后，后面的上下文就更没见过了，再也回不到正轨。这就是数据稀疏在生成时的样子。
  · **长程依赖**：n = 8、10 时，人物介绍里"住在哪个城市"只有约 11% 是对的，和瞎猜（1/8）差不多——
    名字在城市前面超过 9 个字符（"kari lives in " 就有 14 个字符），模型根本看不到是谁；
    n = 14 时上下文刚好盖住名字，正确率升到 60%。评论的评分离食物形容词有 30 多个字符，
    所以即使 n = 14，生成的评论也常常"形容词说好、评分说差"（见 n = 10 的第 2 个样例）。
  · 流畅 ≠ 正确：格式完美但内容是假的句子，是 §5 "可靠性"风险的一个最小例子。
    一个好的语言模型要同时做到三件事：格式对、内容真、不只是照抄。n-gram 很难兼顾，
    代码 02 用同样的指标检查 base model（一个小 Transformer）。""")


if __name__ == "__main__":
    demo_lm_is_a_distribution()
    demo_temperature()
    world = build_world()
    demo_ngram_sparsity(world)
    demo_ngram_samples(world)