"""
=============================================================================
第 11 讲 · 代码 2：把训练过程当成压缩器——在线（prequential）编码长度
=============================================================================

对应讲义：lectures/lecture11_compression.md 第 4 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
代码 01 演示 4 说明：如果要先把模型发过去，大模型压缩小数据是亏的。Jack Rae 的报告用的是另一种编码：
发送方和接收方运行**同一份训练代码**，一边训练一边编码——每一段数据都用"还没见过它"的模型来编码，
接收方解出这段数据之后，用它做同样的一步训练，于是两边的模型始终相同。总码长 = 训练损失曲线下的面积
（再加上训练代码本身的长度，很小）。模型参数一个都不用发。

    演示 1：一个 173 万字符的"更大的迷你互联网"上，4 种宽度的 TinyGPT 各自一边学一边编码：
            在线码长随模型变大而变短；对照 gzip / bz2 / lzma 和"自适应 n-gram"（边数数边编码）。
    演示 2：同样 4 个模型的"两段式"码长（训练完的模型参数 + 用它编码整个数据流）：大模型反而吃亏。
    演示 3：数据流的每十分之一，各个压缩器的边际码率——"学习曲线"就是"压缩曲线"。
    演示 4：Rae 报告里 LLaMA-65B 的账（数字来自对报告的转述），只核对算术。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_prequential_coding.py      # 约 70 秒（4 个小模型各在线训练一遍 + 自适应 n-gram）；需要 PyTorch

依赖 code/lecture01 的 fm_world.py（只用它的生成规则造一个更大的世界，不碰第 1 讲的缓存）、tiny_gpt.py、console.py。

【PyTorch 小知识】
    - 在线训练时，每一段数据**先**算损失（这就是它的码长），**再**用它做一步梯度更新——顺序不能反，
      否则就是用"已经看过答案的模型"来编码，接收方没法复现。
"""

import bz2
import gzip
import lzma
import math
import os
import sys
from collections import Counter, defaultdict

import torch
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))
from console import pad, print_header  # noqa: E402
from fm_world import ALPHABET, build_world  # noqa: E402
from tiny_gpt import CharTokenizer, GPTConfig, count_parameters, make_model  # noqa: E402

B, T = 8, 96                       # 每一步 8 段、每段 96 个字符 → 768 个字符
WIDTHS = [16, 32, 64, 128]
LR = 3e-3


def make_stream():
    """
    一个更大的迷你互联网：和第 1 讲同样的生成规则（同样的 60 个人），评论、故事、算术各 15,000 行（第 1 讲各 2,000 行）。
    返回整段文本（行与行之间用换行隔开）。
    """
    world = build_world(seed=1, n_math=15000, n_reviews=15000, n_stories=15000)
    return world.train_text()


def online_code_length(ids, d):
    """
    在线（prequential）编码：一个宽度为 d 的 2 层 TinyGPT 从随机初始化开始，按顺序读数据流，
    每 768 个字符：先算交叉熵（= 这一段的码长），再用它做一步 AdamW 更新。

    【思路】接收方用同一份代码、同一个随机种子初始化模型；每解出一段，就做同样的一步更新。
            所以发送方不需要发送任何参数——码长就是每一段在"当时的模型"下的交叉熵之和，即训练损失曲线下的面积。
    【参数】ids：整个数据流的字符编号（开头补了一个换行）；d：模型宽度
    【返回】(每一段的 bit 数列表, 训练完的模型)
    【边界】最后不足 768 个字符的零头不编码（只占数据流的十万分之几，忽略）。
    """
    model = make_model(GPTConfig(vocab_size=len(ALPHABET), block_size=T, n_layer=2, n_head=4, d_model=d), seed=0)
    opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.1, betas=(0.9, 0.95))
    C = B * T
    bits = []
    model.train()
    for k in range((len(ids) - 1) // C):
        s = k * C
        x = torch.tensor([ids[s + r * T:s + r * T + T] for r in range(B)])
        y = torch.tensor([ids[s + r * T + 1:s + r * T + T + 1] for r in range(B)])
        for g in opt.param_groups:
            g["lr"] = LR * min(1.0, (k + 1) / 20)                     # 前 20 步热身，之后学习率不变
        loss = F.cross_entropy(model(x).view(-1, len(ALPHABET)), y.view(-1))
        bits.append(loss.item() * C / math.log(2))                   # 先记码长……
        opt.zero_grad()
        loss.backward()                                               # ……再学
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
    return bits, model


@torch.no_grad()
def final_model_bits(model, ids):
    """训练完的模型给整个数据流编码要多少 bit（同样的切法：每段 96 个字符，上下文只在段内）。"""
    model.eval()
    C = B * T
    total = 0.0
    n_chunks = (len(ids) - 1) // C
    for k0 in range(0, n_chunks, 64):
        xs, ys = [], []
        for k in range(k0, min(n_chunks, k0 + 64)):
            s = k * C
            xs += [ids[s + r * T:s + r * T + T] for r in range(B)]
            ys += [ids[s + r * T + 1:s + r * T + T + 1] for r in range(B)]
        logits = model(torch.tensor(xs))
        total += F.cross_entropy(logits.view(-1, len(ALPHABET)), torch.tensor(ys).view(-1), reduction="sum").item()
    return total / math.log(2)


def adaptive_ngram_bits(text, n, k=0.01):
    """
    自适应 n-gram 的在线码长：一边编码一边数数（每个字符先用当前计数编码，再把它计入）。

    【思路】这是经典的"自适应"压缩器（PPM 一类方法的简化版）：和在线训练的神经网络一样，不需要发送任何模型。
            上下文跨行（就是前 n − 1 个字符），开头用换行补齐。加 k 平滑保证概率不为 0。
    【返回】总 bit 数
    """
    V = len(ALPHABET)
    counts = defaultdict(Counter)
    totals = defaultdict(int)
    s = "\n" * (n - 1) + text
    bits = 0.0
    for i in range(n - 1, len(s)):
        h, c = s[i - n + 1:i], s[i]
        bits -= math.log2((counts[h][c] + k) / (totals[h] + k * V))
        counts[h][c] += 1
        totals[h] += 1
    return bits


def bz2_prefix_bits(data, fractions):
    """bz2 压缩数据流的前 f 部分要多少 bit（f 取 fractions 里的每个值）。用来算 bz2 的"边际码率"。"""
    return [8 * len(bz2.compress(data[:int(len(data) * f)], 9)) for f in fractions]


def main():
    stream = make_stream()
    tok = CharTokenizer(ALPHABET)
    ids = [tok.stoi["\n"]] + tok.encode(stream)
    n_chars = ((len(ids) - 1) // (B * T)) * B * T                 # 实际编码的字符数（去掉最后的零头）
    data = stream.encode()

    print_header(f"演示 1：一边学一边编码——{len(stream):,} 个字符的数据流（每步 768 个字符，共 {n_chars // (B * T):,} 步）")
    results = {}
    for d in WIDTHS:
        bits, model = online_code_length(ids, d)
        results[d] = (bits, model)
    print(f"{pad('压缩器', 36)}{pad('参数个数', 12, '>')}{pad('在线码长 bit/字符', 20, '>')}{pad('压缩后', 12, '>')}")
    for d in WIDTHS:
        bits, model = results[d]
        rate = math.fsum(bits) / n_chars
        print(f"{pad(f'TinyGPT，2 层，宽度 {d}（在线训练）', 36)}{pad(f'{count_parameters(model):,}', 12, '>')}"
              f"{pad(f'{rate:.4f}', 20, '>')}{pad(f'{math.fsum(bits) / 8 / 1024:.1f} KB', 12, '>')}")
    for n in (3, 5, 8):
        b = adaptive_ngram_bits(stream, n)
        print(f"{pad(f'自适应 {n}-gram（边数数边编码）', 36)}{pad('—', 12, '>')}{pad(f'{b / len(stream):.4f}', 20, '>')}{pad(f'{b / 8 / 1024:.1f} KB', 12, '>')}")
    for name, fn in [("gzip -9", lambda x: gzip.compress(x, 9, mtime=0)), ("bz2 -9", lambda x: bz2.compress(x, 9)),
                     ("lzma -9", lambda x: lzma.compress(x, preset=9))]:
        c = fn(data)
        print(f"{pad(name, 36)}{pad('—', 12, '>')}{pad(f'{8 * len(c) / len(stream):.4f}', 20, '>')}{pad(f'{len(c) / 1024:.1f} KB', 12, '>')}")
    print(f"（原始数据 {len(data) / 1024:.1f} KB）")

    print_header("演示 2：同样 4 个模型的'两段式'码长：先发训练好的参数，再发用它编码的数据")
    print(f"{pad('宽度', 6)}{pad('参数 × 32 bit', 16, '>')}{pad('数据 | 训练完的模型', 22, '>')}{pad('两段式合计 bit/字符', 22, '>')}{pad('在线码长', 12, '>')}")
    for d in WIDTHS:
        bits, model = results[d]
        p_bits = count_parameters(model) * 32
        data_bits = final_model_bits(model, ids)
        two_part = (p_bits + data_bits) / n_chars
        print(f"{pad(d, 6)}{pad(f'{p_bits / n_chars:.4f}', 16, '>')}{pad(f'{data_bits / n_chars:.4f}', 22, '>')}"
              f"{pad(f'{two_part:.4f}', 22, '>')}{pad(f'{math.fsum(bits) / n_chars:.4f}', 12, '>')}")
    print("（前两列都折算成 bit/字符：参数的代价摊到每个字符上）")

    print_header("演示 3：数据流的每十分之一，各压缩器的边际码率（bit/字符）——学习曲线就是压缩曲线")
    fractions = [i / 10 for i in range(1, 11)]
    bz = bz2_prefix_bits(data, fractions)
    print(f"{pad('数据流', 10)}" + "".join(pad(f"宽度 {d}", 10, ">") for d in WIDTHS) + f"{pad('bz2', 10, '>')}")
    for i, f in enumerate(fractions):
        cells = []
        for d in WIDTHS:
            bits = results[d][0]
            lo, hi = int(len(bits) * (f - 0.1)), int(len(bits) * f)
            cells.append(pad(f"{math.fsum(bits[lo:hi]) / ((hi - lo) * B * T):.3f}", 10, ">"))
        prev = bz[i - 1] if i else 0
        seg_chars = int(len(data) * f) - int(len(data) * (f - 0.1))
        label = f"{int(100 * (f - 0.1))}–{int(100 * f)}%"
        print(f"{pad(label, 10)}" + "".join(cells) + pad(f"{(bz[i] - prev) / seg_chars:.3f}", 10, ">"))

    print_header("演示 4：Rae 报告里的 LLaMA-65B（数字来自对报告的转述，这里只核对算术）")
    tokens, bytes_per_token, nll_bytes, code_bytes = 1.4e12, 4, 400e9, 1e6
    raw = tokens * bytes_per_token
    print(f"训练数据：{tokens:.1e} 个 token × 约 {bytes_per_token} 字节/token = {raw / 1e12:.1f} TB")
    print(f"训练损失曲线下的面积（负对数似然之和）约 400 GB = {nll_bytes * 8 / tokens:.2f} bit/token"
          f"（= {nll_bytes * 8 / tokens * math.log(2):.2f} nats/token，训练全程的平均损失）")
    print(f"训练代码约 1 MB，可以忽略 → 压缩率约 {raw / (nll_bytes + code_bytes):.0f} 倍")
    print(f"如果改用两段式、把 65.2B 个 16 位参数（{65.2e9 * 2 / 1e9:.0f} GB）也发过去：合计 {(nll_bytes + 65.2e9 * 2) / 1e9:.0f} GB，"
          f"压缩率约 {raw / (nll_bytes + 65.2e9 * 2):.1f} 倍")

    print("\n观察 Observations：")
    print("  1. 在线码长随模型变大而单调变短：宽度 16 → 128（参数多 47 倍），1.73 → 0.72 bit/字符。")
    print("     大模型学得更快，损失曲线下的面积更小——这是 Rae 说'更大的模型是更好的压缩器'的意思：不用为参数付钱。")
    print("  2. 两段式编码正好相反：宽度 128 的模型光参数就要摊上 7.66 bit/字符，比原始数据的 8 bit 也差不多了。")
    print("     同一组模型，两种记账给出相反的排序：问'谁压缩得最好'之前，要先说清楚'模型算不算钱'。")
    print("  3. 但在这个玩具世界里，bz2（0.44）比所有在线训练的神经网络和自适应 n-gram（最好 0.60）都好——")
    print("     这个世界的文本由几十种模板生成，信息量本来就很低（第 1 讲的 base model 充分训练后约 0.35 bit/字符），")
    print("     bz2 的块排序变换很擅长这种高度规则、大量相似片段的文本；")
    print("     而神经网络是从零开始学的，前 10% 的数据几乎都花在'学会这是什么语言'上（演示 3 的第一行）。")
    print("     神经网络的边际码率一直在下降，最宽的模型最后 10% 已经到 0.50 左右；bz2 的边际码率则早早稳定。")
    print("     在线学习的压缩器要在'足够长'的数据流上才划算——LLaMA 读了 1.4 万亿个 token，我们只有 173 万个字符。")
    print("  4. LLaMA 的账：400 GB 的面积相当于平均 2.29 bit/token（1.58 nats），压缩率约 14 倍；")
    print("     就算把 130 GB 的参数也算上，仍有约 10.6 倍——在 5.6 TB 的规模上，参数的代价已经不重要了。")


if __name__ == "__main__":
    main()