"""
=============================================================================
第 4 讲 · 代码 2：一个"知道熵"的合成语言，以及在它上面的扫描训练（模型大小 × 数据量）
=============================================================================

对应讲义：lectures/lecture04_scaling_laws.md 第 4 节（以及第 5、6 节用到的数据）

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
规模定律是从很多次训练里"量"出来的：换模型大小、换数据量，每次记下最后的损失。
在真实语料上，损失最低能到多少（"不可约损失" E）谁也不知道。这里换一个熵已知的合成语言，
在上面训练 6 种宽度 × 4 种数据量的小 Transformer，看损失怎样随 N 和 D 变化、怎样趋向 E。

    演示 1：合成语言长什么样；它的熵率 E 怎样精确算出来；"真实模型"在验证文本上确实做到 E；
            n-gram（数数）能走多远。
    演示 2：扫描训练（6 个宽度 × 4 个数据量 = 24 次训练，外加 2 次换种子的重复）。第一次运行约 2 分钟，
            结果缓存到 cache/sweep.json，之后几秒钟读完。打印损失表和"损失 − E"表。
    演示 3：两个切片——固定 D 看 L 随 N 怎样变，固定 N 看 L 随 D 怎样变；对数坐标下的局部斜率。
    演示 4：噪声有多大——同一个设定换一个初始化种子，损失差多少。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_toy_language_sweep.py           # 第一次约 2 分钟（训练）；有缓存时约 5 秒
    python3 02_toy_language_sweep.py --retrain # 删掉缓存、重新训练（结果逐字相同）

需要 numpy、PyTorch；依赖本目录的 scaling_tools.py 和第 1 讲的 tiny_gpt.py、console.py（只读）。
代码 03、04 读取这里的缓存；如果没有缓存，它们会自己先训练（同样约 2 分钟）。
"""

import math
import os
import sys
import time

import numpy as np

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 scaling_tools import (ALPHABET_L4, CACHE_PATH, DATA_SIZES, LANG_CONFIG, TRAIN_CONFIG, WIDTHS,  # noqa: E402
                           build_language, ensure_runs, grid_specs, grid_table, make_stream)


# =============================================================================
# 演示 1：合成语言与它的熵
# =============================================================================
def ngram_bits(train, val, n, k=0.01):
    """
    字符级 n-gram 在验证文本上的交叉熵（bit/字符），加 k 平滑。

    【思路】把"前 n−1 个字符"编码成一个整数（21 进制），用 np.bincount 数出每个（上下文, 下一个字符）出现几次——
            和第 1 讲 ngram_lm.py 做的事一样，只是用 numpy 一次性数完几百万个字符。
    【参数】train、val：一维 numpy 整数数组（字符编号）；n：阶数；k：平滑常数
    【返回】(交叉熵, 计数表里非零项的个数)——后者可以粗略地看成 n-gram 模型的"参数量"
    【边界】n 太大时上下文种类 21^(n−1) 会爆炸，这里只用到 n = 5（约 400 万种上下文）。
    【numpy 小知识】np.bincount(x, minlength=m) 返回长度 m 的数组，第 i 个元素是 i 在 x 里出现的次数。
    """
    V = len(ALPHABET_L4)

    def codes(a):
        # 第 i 个 n-gram 的编码：Σ a[i+j] · V^(n−1−j)
        c = np.zeros(len(a) - n + 1, dtype=np.int64)
        for j in range(n):
            c = c * V + a[j:len(a) - n + 1 + j]
        return c

    tr, va = codes(train), codes(val)
    joint = np.bincount(tr, minlength=V ** n).astype(np.float64)
    ctx = joint.reshape(-1, V).sum(axis=1)
    p = (joint[va] + k) / (ctx[va // V] + k * V)
    return float(-np.log2(p).mean()), int((joint > 0).sum())


def demo_language(lang):
    print_header("演示 1：一个熵可以精确计算的合成语言")
    ids = lang.sample_word_ids(40, seed=123)
    print("采样一段文本（词由辅音 + 元音的音节拼成，词之间用空格隔开）：")
    print("   " + lang.text(ids)[:96])
    print("   " + lang.text(ids)[96:192])
    order = np.argsort(-lang.pi)
    top = ", ".join(f"{lang.words[i]} {lang.pi[i]:.3f}" for i in order[:5])
    rare = ", ".join(f"{lang.words[i]} {lang.pi[i]:.5f}" for i in order[-3:])
    print(f"\n词表 {LANG_CONFIG['n_words']} 个词。最常见的 5 个词（长期出现比例）：{top}")
    print(f"最罕见的 3 个：{rare}")
    i0 = int(order[0])
    favs = ", ".join(lang.words[j] for j in lang.fav[i0])
    print(f"搭配：'{lang.words[i0]}' 后面有 {LANG_CONFIG['gamma']:.0%} 的概率接它偏爱的两个词之一（{favs}），"
          f"其余按背景 Zipf 分布抽。")

    print("\n熵率的计算（全部来自生成过程的参数，不需要任何数据）：")
    print(f"   每个词的条件熵 Σ π_i H(P(·|i))  = {lang.bits_per_word:.4f} bit/词")
    print(f"   每个词平均占的字符数（含空格）    = {lang.chars_per_word:.4f}")
    print(f"   熵率 E = 两者相除               = {lang.entropy_rate:.4f} bit/字符   ← 任何模型的损失下限")

    train = make_stream(lang, DATA_SIZES[-1], TRAIN_CONFIG["train_seed"]).numpy().astype(np.int64)
    val_ids = lang.sample_word_ids(int(TRAIN_CONFIG["eval_chars"] / lang.chars_per_word), TRAIN_CONFIG["val_seed"])
    val = make_stream(lang, TRAIN_CONFIG["eval_chars"], TRAIN_CONFIG["val_seed"]).numpy().astype(np.int64)
    print(f"\n'真实模型'（知道生成过程）在验证文本上的平均编码长度：{lang.oracle_bits(val_ids):.4f} bit/字符"
          f"（和 E 只差抽样噪声）")

    print(f"\n数数能走多远？在 {len(train) / 1e6:.0f}M 个训练字符上数 n-gram，在验证文本上算交叉熵：")
    rows = [("均匀分布（什么都没学）", math.log2(len(ALPHABET_L4)), None)]
    for n in range(1, 6):
        bits, nz = ngram_bits(train, val, n)
        rows.append((f"{n}-gram", bits, nz))
    rows.append(("熵率 E（下限）", lang.entropy_rate, None))
    for name, bits, nz in rows:
        extra = f"   计数表里 {nz:,} 个非零项" if nz else ""
        print(f"   {pad(name, 26)} {bits:.4f} bit/字符{extra}")
    drops = [rows[i][1] - rows[i + 1][1] for i in range(1, 5)]
    print("""
观察 Observations：
  · 这个语言的规律分三层：字母层（辅音、元音交替）、词层（300 个词怎么拼）、词与词之间（搭配）。
    n-gram 每加一阶降得越来越少（{d0:.2f}、{d1:.2f}、{d2:.2f}、{d3:.2f}）：1-gram → 2-gram 学会了辅音元音交替，
    之后要一个字母一个字母地"认出"词。5-gram 只能看到 4 个字符，看不全一个长词，更看不到上一个词。
  · 注意 5-gram 的 {b5:.3f} 并不差——下面的扫描里，只有最宽的两个 Transformer 在 2M 数据上超过了它，
    而它的计数表只有 {nz5:,} 个非零项。原因是我们的 Transformer 训练得很短（最多约 2000 步），离收敛还远；
    n-gram 只要把数据数一遍。但 5-gram 看不到 4 个字符以外，数据再多也到不了 E；Transformer 原则上可以。
    这提醒我们：规模定律描述的永远是"某一种模型 + 某一种训练方法"下损失怎样随规模变化。
  · 要做到 E，模型需要记住整个词表、每个词偏爱的后继词，并且能"回头看"十几个字符——
    越大的模型能记住的越多，越多的数据能让它见到越多罕见的词和搭配。这正是规模定律要量化的东西。
  · E 是**精确**的：它从生成过程直接算出来，"真实模型"在验证文本上的编码长度也确实落在 E 附近。
    真实语料没有这样的真值，规模定律里的 E 只能拟合出来（代码 03 检查拟合出的 E 靠不靠谱）。""".format(
        d0=drops[0], d1=drops[1], d2=drops[2], d3=drops[3], b5=rows[5][1], nz5=rows[5][2]))


# =============================================================================
# 演示 2：扫描训练
# =============================================================================
def demo_sweep(lang):
    print_header("演示 2：扫描训练 —— 6 个宽度 × 4 个数据量（第一次运行约 2 分钟，之后读缓存）")
    print(f"所有模型都是 2 层 Transformer（第 1 讲的 TinyGPT），只改宽度 d_model；每个注意力头 8 维。")
    print(f"每步 {TRAIN_CONFIG['batch_size']} × {TRAIN_CONFIG['block_size']} = "
          f"{TRAIN_CONFIG['batch_size'] * TRAIN_CONFIG['block_size']} 个 token；数据按顺序读、每个 token 只读一遍；"
          f"学习率调度的长度 = 训练长度（余弦衰减）。\n")
    runs = ensure_runs(grid_specs())
    table, sizes = grid_table(runs)
    E = lang.entropy_rate

    head = pad("宽度 d", 8) + pad("总参数 N", 10) + pad("非嵌入参数", 12) + "".join(
        pad(f"D={D / 1e6:g}M", 10) for D in DATA_SIZES)
    print("验证损失（bit/字符）：")
    print("  " + head)
    for d in WIDTHS:
        n_tot, n_ne = sizes[d]
        print("  " + pad(d, 8) + pad(f"{n_tot:,}", 10) + pad(f"{n_ne:,}", 12)
              + "".join(pad(f"{table[d][D]:.4f}", 10) for D in DATA_SIZES))
    print(f"  （熵率 E = {E:.4f}）")
    print("\n'还没学会的部分' L − E：")
    print("  " + pad("宽度 d", 8) + "".join(pad(f"D={D / 1e6:g}M", 10) for D in DATA_SIZES))
    for d in WIDTHS:
        print("  " + pad(d, 8) + "".join(pad(f"{table[d][D] - E:.4f}", 10) for D in DATA_SIZES))
    best = min((table[d][D], d, D) for d in WIDTHS for D in DATA_SIZES)
    print(f"\n最好的一格：d = {best[1]}、D = {best[2] / 1e6:g}M，损失 {best[0]:.4f}，离 E 还差 {best[0] - E:.4f}。")
    c_small = 6 * sizes[WIDTHS[0]][0] * DATA_SIZES[0]
    c_big = 6 * sizes[WIDTHS[-1]][0] * DATA_SIZES[-1]
    print(f"算力 C = 6ND 从 {c_small:.1e} 到 {c_big:.1e} FLOPs（跨 {math.log10(c_big / c_small):.1f} 个数量级）；"
          f"作为对照，GPT-3 约 3.1e23。")
    D0, D1, D2, D3 = DATA_SIZES
    lo, hi = WIDTHS[0], WIDTHS[-1]
    shrink_lo = 1 - (table[lo][D3] - E) / (table[lo][D2] - E)
    shrink_hi = 1 - (table[hi][D3] - E) / (table[hi][D2] - E)
    min_excess_small = min(table[d][D0] for d in WIDTHS) - E
    print(f"""
观察 Observations：
  · 同一行里（固定 N）数据越多越好，没有例外。同一列里（固定 D），D ≥ 1M 时越宽越好；
    D 小的两列有例外：d = 32 在 0.25M、0.5M 上不如 d = 24。
  · D = 0.25M（只有 244 步）那一列：所有模型都离 E 很远（L − E 最小也有 {min_excess_small:.2f}），排序也不规则——
    步数太少，训练还在剧烈变化，谁先"开窍"有很大的偶然性（演示 4）。这是 Hestness et al. 说的"小数据区"
    （small data region）：拟合规模定律时要排除这样的点，或者保证每个点都训练得足够（代码 03、04 只用 D ≥ 0.5M）。
  · 一个资源的回报取决于另一个资源：从 1M 到 2M，最窄的模型"还没学会的部分"只减少了 {shrink_lo:.0%}
    （{table[lo][D2] - E:.3f} → {table[lo][D3] - E:.3f}），容量（N）成了瓶颈；最宽的模型减少了 {shrink_hi:.0%}
    （{table[hi][D2] - E:.3f} → {table[hi][D3] - E:.3f}）。Kaplan 的 L(N, D) 公式和 Chinchilla 的 E + A/N^α + B/D^β
    描述的就是这种"两个瓶颈此消彼长"。
  · 即使最好的一格，离 E 也还有 {best[0] - E:.2f}：我们的模型和数据都太小，还在"能学的还很多"的区域。""")
    return runs, table, sizes


# =============================================================================
# 演示 3：两个切片与局部斜率
# =============================================================================
def local_slopes(xs, ys):
    """相邻两点之间 log(y) 对 log(x) 的斜率（取负号，所以下降时为正数）。"""
    return [-(math.log(ys[i + 1]) - math.log(ys[i])) / (math.log(xs[i + 1]) - math.log(xs[i]))
            for i in range(len(xs) - 1)]


def demo_slices(lang, table, sizes):
    print_header("演示 3：两个切片 —— 固定 D 看 N，固定 N 看 D（对数坐标下的局部斜率）")
    E = lang.entropy_rate
    D = DATA_SIZES[-1]
    Ns = [sizes[d][0] for d in WIDTHS]
    Ls = [table[d][D] for d in WIDTHS]
    s_raw = local_slopes(Ns, Ls)
    s_ex = local_slopes(Ns, [l - E for l in Ls])
    print(f"(a) 固定 D = {D / 1e6:g}M，改变 N：")
    print("  " + pad("N", 10) + pad("L", 9) + pad("L − E", 9) + pad("L 的局部斜率", 15) + "L − E 的局部斜率")
    for i, d in enumerate(WIDTHS):
        a = f"{s_raw[i]:.3f}" if i < len(s_raw) else ""
        b = f"{s_ex[i]:.3f}" if i < len(s_ex) else ""
        print("  " + pad(f"{Ns[i]:,}", 10) + pad(f"{Ls[i]:.4f}", 9) + pad(f"{Ls[i] - E:.4f}", 9) + pad(a, 15) + b)
    d = WIDTHS[-1]
    Ds = list(DATA_SIZES)
    Ld = [table[d][x] for x in Ds]
    t_raw = local_slopes(Ds, Ld)
    t_ex = local_slopes(Ds, [l - E for l in Ld])
    print(f"\n(b) 固定 N = {sizes[d][0]:,}（d = {d}），改变 D：")
    print("  " + pad("D", 10) + pad("L", 9) + pad("L − E", 9) + pad("L 的局部斜率", 15) + "L − E 的局部斜率")
    for i, x in enumerate(Ds):
        a = f"{t_raw[i]:.3f}" if i < len(t_raw) else ""
        b = f"{t_ex[i]:.3f}" if i < len(t_ex) else ""
        print("  " + pad(f"{x / 1e6:g}M", 10) + pad(f"{Ld[i]:.4f}", 9) + pad(f"{Ld[i] - E:.4f}", 9) + pad(a, 15) + b)
    print("""
（"局部斜率"是相邻两点连线在 log-log 坐标下的斜率。纯幂律 y = A·x^(−α) 的局部斜率处处等于 α。）

观察 Observations：
  · (a) 直接看 L，局部斜率从 {r0:.2f} 降到 {r1:.2f}——曲线在对数坐标下是弯的，因为 L 里有一块减不掉的 E。
    减掉 E 以后，斜率在 {e_lo:.2f}–{e_hi:.2f} 之间跳动、没有一路变小的趋势：弯曲主要来自 E。
    跳动本身是噪声（演示 4）：相邻两个宽度的 N 只差 2 倍左右，损失的一点随机波动就能让局部斜率变很多。
    所以没有人用"两点连线"估指数，而是用很多点一起拟合（代码 03）。
  · (b) 数据方向的局部斜率（{t0:.2f} → {t1:.2f}）比参数方向大，而且一路变小：固定 N 时，D 越大，
    "还没学会的部分"里"容量不够"的那一块占的比例越大，L 趋向的不是 E，而是"E + 容量不够造成的损失"。
  · 两个方向的指数都比真实 LLM 大得多（Kaplan：α_N ≈ 0.076、α_D ≈ 0.095）。我们的语言只有 300 个词，
    "能学的东西"有限，学得也快；真实语言的长尾长得多。这是玩具和真实世界在"量级"上的差别，机制是一样的。""".format(
        r0=s_raw[0], r1=s_raw[-1], e_lo=min(s_ex), e_hi=max(s_ex), t0=t_ex[0], t1=t_ex[-1]))


# =============================================================================
# 演示 4：噪声
# =============================================================================
def demo_noise(runs):
    print_header("演示 4：噪声 —— 同一个设定换一个初始化种子")
    print("  " + pad("设定", 22) + pad("种子 0", 10) + pad("种子 1", 10) + "差")
    diffs = []
    for d in (16, 48):
        a = runs[f"L2_d{d}_D500000_S500000_seed0"]["loss"]
        b = runs[f"L2_d{d}_D500000_S500000_seed1"]["loss"]
        diffs.append(abs(a - b))
        print("  " + pad(f"d = {d}，D = 0.5M", 22) + pad(f"{a:.4f}", 10) + pad(f"{b:.4f}", 10) + f"{b - a:+.4f}")
    print(f"""
观察 Observations：
  · 只换初始化种子，损失就差 {min(diffs):.3f}–{max(diffs):.3f} bit/字符。拟合规模定律时，比这个还小的差别不能当真。
  · Kaplan et al. 估计他们的训练"换种子"的波动约为 0.02（nats/token），并用它来定义"过拟合开始变得显著"的阈值。
    大模型、大数据上的噪声相对小得多；我们的小模型、短训练噪声更大，这是玩具实验必须承认的局限。""")


def main():
    if "--retrain" in sys.argv and os.path.exists(CACHE_PATH):
        os.remove(CACHE_PATH)
        print("（已删除缓存，重新训练）")
    t0 = time.time()
    lang = build_language()
    demo_language(lang)
    runs, table, sizes = demo_sweep(lang)
    demo_slices(lang, table, sizes)
    demo_noise(runs)
    print(f"\n（总用时 {time.time() - t0:.0f} 秒）")


if __name__ == "__main__":
    main()