"""
=============================================================================
第 11 讲 · 代码 3：最小描述长度（MDL）选模型，以及"压缩 = 智能"的边界
=============================================================================

对应讲义：lectures/lecture11_compression.md 第 5、6 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
    演示 1：用码长选模型。n-gram 的阶数 n 从 1 到 10：训练集损失总是选最大的 n（背得越多越好），
            验证集损失在中间某处最好；在线码长（不看验证集！）和两段式码长分别选哪个 n？换三种训练集大小看一看。
    演示 2：好的压缩器只对"它的世界"好：把验证集倒过来写、打乱字符、换成随机字符，
            base model 和 bz2 各要多少 bit/字符。
    演示 3：重复的随机串：同一个随机串在一行里出现两次，第二次能便宜多少？base model、自适应 8-gram、gzip 各自怎样——
            这是第 12 讲 induction head 的伏笔。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 03_mdl_and_limits.py       # 约 15 秒；需要 PyTorch（演示 2、3 用到 base model）

依赖本目录的 codec_tools.py，以及 code/lecture01 的 fm_world.py、ngram_lm.py、tiny_gpt.py、console.py（base model 只读加载）。

【Python 小知识】
    - min(range(…), key=函数)：返回让函数值最小的那个下标——"选使码长最短的 n"就是一行。
    - random.Random(seed).shuffle(列表)：原地打乱列表。
"""

import bz2
import gzip
import lzma
import math
import os
import random
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))
from codec_tools import gpt_distributions, ideal_bits  # noqa: E402
from console import pad, print_header  # noqa: E402
from fm_world import ALPHABET, build_world  # noqa: E402
from ngram_lm import CharNGramLM  # noqa: E402
from tiny_gpt import CharTokenizer, load_or_train_base_model  # noqa: E402

V = len(ALPHABET)
ORDERS = list(range(1, 11))


def prefix_lines(lines, n_chars):
    """取前若干行，使总字符数刚刚达到 n_chars（行与行之间用换行隔开，末尾也有换行）。"""
    out, c = [], 0
    for line in lines:
        out.append(line)
        c += len(line) + 1
        if c >= n_chars:
            break
    return "\n".join(out) + "\n"


def prequential_bits(text, n, k=0.01):
    """
    自适应 n-gram 在 text 上的在线码长（bit）：每个字符先用"目前为止数到的计数"编码，再把它计入。

    【思路】和第 1 讲 CharNGramLM 完全相同的模型族（每行开头补 n − 1 个换行、加 k 平滑），只是一边编码一边训练。
            它不需要看验证集，也不需要为模型付钱——但大 n 的模型在开头学得慢（大部分上下文都还没见过），
            要为此付出更长的码长。
    """
    lm = CharNGramLM(n, ALPHABET, k=k)
    bits = 0.0
    for h, c in lm._events(text):
        bits -= math.log2(lm.prob(h, c))
        lm.counts[h][c] += 1
        lm.totals[h] += 1
    return bits


def table_bits(lm):
    """
    两段式编码里"发送 n-gram 计数表"要多少 bit：每个非零计数记一条（上下文 + 字符，共 n 个字符，各 log2 44 bit；
    计数本身用 Elias γ 码，约 2·log2(计数 + 1) + 1 bit）。

    【边界】这是一种朴素的写法：真实的模型编码可以聪明得多（共享前缀、用概率模型编码计数……），
            码长会小很多。MDL 的结论依赖于"模型怎样编码"——演示 1 会看到这一点。
    """
    return math.fsum(lm.n * math.log2(V) + 2 * math.log2(cnt + 1) + 1
                     for cn in lm.counts.values() for cnt in cn.values())


def demo1_mdl(world):
    print_header("演示 1：用码长选 n-gram 的阶数 n（1–10）")
    val = world.val_text()
    summary = []
    for size in (10_000, 30_000, None):
        train = prefix_lines(world.train_lines, size) if size else world.train_text()
        rows = []
        for n in ORDERS:
            lm = CharNGramLM(n, ALPHABET, k=0.01).fit(train)
            tr = lm.bits_per_char(train)
            rows.append(dict(n=n, train=tr, val=lm.bits_per_char(val),
                             preq=prequential_bits(train, n) / len(train),
                             two=(table_bits(lm) + tr * len(train)) / len(train)))
        if size == 30_000:
            print(f"训练集 {len(train):,} 个字符时的完整表格（单位 bit/字符）：")
            print(f"{pad('n', 4)}{pad('训练集', 10, '>')}{pad('验证集', 10, '>')}{pad('在线码长', 10, '>')}{pad('两段式', 10, '>')}")
            for r in rows:
                print(f"{pad(r['n'], 4)}" + "".join(pad(f"{r[key]:.3f}", 10, ">") for key in ("train", "val", "preq", "two")))
            print()
        pick = {key: min(rows, key=lambda r: r[key])["n"] for key in ("train", "val", "preq", "two")}
        summary.append((len(train), pick, rows))
    print(f"{pad('训练集字符数', 14)}{pad('按训练集损失', 14, '>')}{pad('按验证集损失', 14, '>')}{pad('按在线码长', 12, '>')}{pad('按两段式码长', 14, '>')}")
    for size, pick, _ in summary:
        print(f"{pad(f'{size:,}', 14)}{pad(pick['train'], 14, '>')}{pad(pick['val'], 14, '>')}{pad(pick['preq'], 12, '>')}{pad(pick['two'], 14, '>')}")
    print("（每格是使该列最小的 n）")
    print("\n观察 Observations：")
    print("  1. 训练集损失总是偏爱大 n：上下文越长，越能把训练集背下来。它不能用来选模型。")
    print("  2. 验证集损失在中间最好：训练集只有 1 万个字符时 n = 5，3 万个时 n = 7（n = 8 几乎一样），")
    print("     用全部 26 万个字符时到 n = 10 还在下降——数据越多，越复杂的模型越划算。")
    print("  3. 在线码长**完全不看验证集**，只问'一边学一边编码，总共要多少 bit'。它在 1 万个字符时正好选中 n = 5，")
    print("     3 万个时选 n = 5（验证集最好是 7），全部数据时选 n = 8：偏保守，但跟着数据量一起变大，方向是对的。")
    print("     原因：大 n 的模型在开头学得慢（大部分上下文都没见过），这一段'学习的代价'被计进了码长——这就是奥卡姆剃刀。")
    print("  4. 两段式码长（朴素地把计数表也发过去）选的 n 太小（2–4）：我们的计数表编码很浪费，模型的代价被高估了。")
    print("     MDL 的结论依赖'模型怎样编码'；在线编码回避了这个问题，这也是 Rae 选择它的原因之一。")


def demo2_limits(world, tok, base):
    print_header("演示 2：好的压缩器只对'它的世界'好（bit/字符；均匀分布是 5.46）")
    lines = world.val_lines
    r1, r2 = random.Random(1), random.Random(2)

    def shuffled(line):
        chars = list(line)
        r1.shuffle(chars)
        return "".join(chars)

    val = world.val_text()
    letters = ALPHABET[1:]                                  # 去掉换行符
    rand = "".join("\n" if i % 60 == 0 else r2.choice(letters) for i in range(1, len(val) + 1))
    variants = [("原样（迷你互联网的验证集）", val),
                ("每一行倒过来写", "\n".join(line[::-1] for line in lines) + "\n"),
                ("每一行内部打乱字符", "\n".join(shuffled(line) for line in lines) + "\n"),
                ("随机字符（每 60 个一行）", rand)]
    print(f"{pad('数据', 28)}{pad('base model', 12, '>')}{pad('gzip', 10, '>')}{pad('bz2', 10, '>')}{pad('lzma', 10, '>')}")
    for name, text in variants:
        g = ideal_bits(text, ALPHABET, gpt_distributions(base, tok, text)) / len(text)
        data = text.encode()
        cells = [8 * len(f(data)) / len(text) for f in (lambda b: gzip.compress(b, 9, mtime=0),
                                                         lambda b: bz2.compress(b, 9), lambda b: lzma.compress(b, preset=9))]
        print(f"{pad(name, 28)}{pad(f'{g:.3f}', 12, '>')}" + "".join(pad(f"{c:.3f}", 10, ">") for c in cells))
    print("\n观察 Observations：")
    print("  1. 在它熟悉的世界里，base model 是最好的压缩器（0.355，bz2 的 1/3）。")
    print("  2. 把同样的文本倒过来写，信息量一点没变，bz2 几乎不受影响（1.169），base model 却要 10.7 bit/字符——")
    print("     比什么都不知道的均匀分布（5.46）还差一倍。它'自信地猜错'：它的预测建立在'迷你互联网是正着写的'这个先验上。")
    print("  3. 随机字符没有任何规律，没有压缩器能低于 5.4 左右（bz2 5.58）；base model 要 15.4 bit/字符，代价最大。")
    print("  4. 所以'压缩 = 智能'的说法要加一个前提：压缩的是**我们关心的那个世界**的数据。一个只见过迷你互联网的模型")
    print("     是迷你互联网的好压缩器、别的东西的坏压缩器。大语言模型能压缩图像和音频（Delétang et al. 2023），")
    print("     靠的是训练数据极其广泛，加上上下文学习——在上下文里临时'认出'新数据的规律。")


def demo3_repeats(tok, base):
    print_header("演示 3：同一个随机串在一行里出现两次——第二次能便宜多少？")
    rng = random.Random(0)
    letters = "abcdefghijklmnopqrstuvwxyz"
    pairs = []
    for _ in range(200):
        s = "".join(rng.choice(letters) for _ in range(20))
        pairs.append(s)
    text = "\n".join(f"{s} {s}" for s in pairs) + "\n"
    first_only = "\n".join(f"{s} " for s in pairs) + "\n"
    idx = {c: i for i, c in enumerate(ALPHABET)}
    d = gpt_distributions(base, tok, text)
    lm = CharNGramLM(8, ALPHABET, k=0.01)
    costs = {"base model": [0.0, 0.0], "自适应 8-gram": [0.0, 0.0]}
    pos = 0
    for s in pairs:
        line = f"{s} {s}\n"
        for j, (h, ch) in enumerate(lm._events(line)):
            part = 0 if j < 20 else (1 if 21 <= j < 41 else None)
            b_gpt = -math.log2(d[pos][idx[ch]])
            b_ng = -math.log2(lm.prob(h, ch))
            lm.counts[h][ch] += 1                         # 自适应：编码之后立刻计入（在线）
            lm.totals[h] += 1
            if part is not None:
                costs["base model"][part] += b_gpt
                costs["自适应 8-gram"][part] += b_ng
            pos += 1
    n_chars = 20 * len(pairs)
    print(f"{pad('压缩器', 18)}{pad('第一次出现', 14, '>')}{pad('第二次出现', 14, '>')}  （bit/字符；随机串的熵是 log2 26 = 4.70）")
    for name, (a, b) in costs.items():
        print(f"{pad(name, 18)}{pad(f'{a / n_chars:.3f}', 14, '>')}{pad(f'{b / n_chars:.3f}', 14, '>')}")
    for name, fn in [("gzip", lambda x: gzip.compress(x, 9, mtime=0)), ("lzma", lambda x: lzma.compress(x, preset=9))]:
        whole, half = fn(text.encode()), fn(first_only.encode())
        extra = 8 * (len(whole) - len(half)) / n_chars
        print(f"{pad(name, 18)}{pad(f'{8 * len(half) / (21 * len(pairs) + len(pairs)):.3f}', 14, '>')}{pad(f'{extra:.3f}', 14, '>')}"
              f"  （第二列 = 只有第一次时的大小；第三列 = 加上第二次之后多出来的部分）")
    print("\n观察 Observations：")
    print("  1. 第一次出现的随机串谁也压不了：gzip、lzma 约 5.0 bit/字符（接近熵 4.70，多出来的是空格、换行和文件头），")
    print("     自适应 8-gram 5.6，base model 要 13.5——它没见过这种文本，而且很自信地猜错。")
    print("  2. 第二次出现时，gzip、lzma 只要 0.4–0.8 bit/字符：它们的核心机制（LZ77）就是'这一段在前面出现过，回去抄'。")
    print("     自适应 8-gram 降到 2.2：第一次出现时它已经把每个 8 字符的片段计了数，第二次就能接着往下预测（开头几个字符还得猜）。")
    print("  3. base model 第二次出现的代价和第一次几乎一样（13.8）：它不会'回去抄'。第 5 讲已经发现，")
    print("     这个模型在训练中没有形成 induction head——'找到前面出现过的同一段，照着往下写'的注意力电路。")
    print("     第 12 讲会专门训练一个形成了 induction head 的小模型，看它怎样获得这种'复制'的能力。")
    print("     从压缩的角度看，induction head 就是神经网络版的 LZ77；它也被认为是上下文学习的一个重要机制。")


if __name__ == "__main__":
    world, tok = build_world(), CharTokenizer(ALPHABET)
    base, _ = load_or_train_base_model(world, tok)
    demo1_mdl(world)
    demo2_limits(world, tok, base)
    demo3_repeats(tok, base)