"""
scaling.py —— 第 12 讲：一个"真实熵已知"的语料、一族大小不同的 Transformer，和规模定律的拟合工具
=====================================================================

【这个文件做什么】
    规模定律（scaling laws）研究的是：测试损失 L 怎样随参数量 N、训练 token 数 D、算力 C ≈ 6ND 变化。
    真实研究要训练几十到几百个模型（Kaplan et al. 2020 最大到 15 亿参数，Chinchilla 训练了 400 多个）。
    这里在笔记本 CPU 上做同样的事，模型从几千到几十万个参数，训练 token 从几万到几百万。

    语料（"花园年鉴"）是两种文档的混合，每篇独立采样、取之不尽（没有重复，除非特意重复）：
        花园句子（40%）：第 1 讲的花园世界，例如 "she picks her ripe plum . <eos>"——有语法、有一致关系，
                        小模型很快就能学会；
        百科条目（60%）：3000 位人物，名字由 3 个音节组成（例如 "ka lo mi"），每篇说一条事实：
                        "ka lo mi lives in Ashby . <eos>"、"ka lo mi grows tulip . <eos>"、"ka lo mi keeps goat . <eos>"。
                        人物被提到的频率服从 Zipf 定律（指数 1.0）：头部的人物家喻户晓，尾部的几乎没人提。
    为什么这样设计：
        (1) **真实熵可以精确算出来**：花园句子有第 1 讲的真实概率；百科条目的每个 token 的真实条件概率
            可以从 Zipf 权重和事实表直接算出。于是不可约损失 E（"语言本身的熵"）是已知的——
            真实研究只能靠拟合去估计它，这里可以检验拟合得准不准；
        (2) **长尾**：要记住的事实（9000 条）和人物的频率都服从 Zipf 分布。看到更多数据就多认识一些冷门人物，
            模型更大就能多记住一些事实——这正是对"为什么损失随规模呈幂律下降"的一种流行解释
            （Hutter 2021；Michaud et al. 2023 的"量子化模型"：能力被拆成很多小块，它们的使用频率服从幂律）。

    这里提供：语料与真实熵、模型族、可缓存的训练（结果写进 results.json）、幂律与 Chinchilla 参数化拟合。

【LLM 联系】
    Kaplan et al. (2020) 在 WebText2 上拟合 L(N)、L(D)、L(C)；Hoffmann et al. (2022, Chinchilla) 用三种方法
    回答"给定算力，N 和 D 怎么分"。本讲的 01～03 号脚本依次复现这两篇论文的主要实验。
"""

import hashlib
import json
import math
import os
import random
import sys
import time

import numpy as np
import torch
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture01", "lecture02"):
    sys.path.insert(0, os.path.join(HERE, "..", _d))

from toy_world import NOUNS, GardenWorld, Vocab, all_words, zipf_weights  # noqa: E402

RESULTS = os.path.join(HERE, "results.json")
CKPT_DIR = os.path.join(HERE, "checkpoints")

# =====================================================================
# 1. 语料
# =====================================================================

SYLLABLES = ["ka", "lo", "mi", "nu", "pe", "ri", "so", "ta", "vu", "xa", "zi", "bo", "du", "fe", "gi", "ho"]
TOWNS = ["Ashby", "Brookfield", "Coveton", "Dunmore", "Elmwick", "Fallowby", "Greystone", "Hollins",
         "Ivydale", "Kestrel", "Larkhill", "Millbrook", "Northam", "Oakridge", "Pebworth", "Redcliff"]
TOWN_PRIOR = [1 / (i + 2) for i in range(len(TOWNS))]
PLANTS = NOUNS["fruit"] + NOUNS["veg"] + NOUNS["flower"]
ANIMALS = NOUNS["animal"]
RELATIONS = ["town", "plant", "animal"]
VERB = {"town": ["lives", "in"], "plant": ["grows"], "animal": ["keeps"]}
N_PEOPLE = 3000
ZIPF_S = 1.0
P_BIO = 0.6

VOCAB = Vocab(["<bos>", "<eos>", "<unk>"] + all_words() + SYLLABLES + TOWNS
              + [w for w in ["lives", "in", "grows", "keeps"] if w not in all_words()])
BOS, EOS = VOCAB.stoi["<bos>"], VOCAB.stoi["<eos>"]


class Almanac:
    """
    "花园年鉴"的生成过程和真实概率。

    【属性】people：[(名字的 3 个音节, {关系: 值})]，按流行度从高到低；weights：每人被提到的概率（Zipf）。
    【方法】
        sample_doc(rng)：采样一篇文档（词列表，以 <eos> 结尾）；
        logprob_steps(doc)：这篇文档每个 token 的真实条件概率的对数（列表）；
        kind(doc)：文档是 "garden" 还是 "bio"。
    """

    def __init__(self, seed=0):
        rng = random.Random(seed)
        names = rng.sample([(a, b, c) for a in SYLLABLES for b in SYLLABLES for c in SYLLABLES], N_PEOPLE)
        self.people = []
        for nm in names:
            facts = {"town": rng.choices(TOWNS, weights=TOWN_PRIOR)[0], "plant": rng.choice(PLANTS),
                     "animal": rng.choice(ANIMALS)}
            self.people.append((nm, facts))
        self.weights = zipf_weights(N_PEOPLE, ZIPF_S)
        self.garden = GardenWorld()
        # 名字前缀 → 概率质量（用来算名字每个音节的真实条件概率）
        self.prefix_mass = {}
        for (nm, _), w in zip(self.people, self.weights):
            for k in range(4):
                self.prefix_mass[nm[:k]] = self.prefix_mass.get(nm[:k], 0.0) + w
        self.by_name = {nm: f for nm, f in self.people}

    def sample_doc(self, rng):
        if rng.random() < P_BIO:
            i = rng.choices(range(N_PEOPLE), weights=self.weights)[0]
            nm, facts = self.people[i]
            r = rng.choice(RELATIONS)
            return list(nm) + VERB[r] + [facts[r], ".", "<eos>"]
        return self.garden.sample_doc(rng)

    @staticmethod
    def kind(doc):
        return "bio" if doc[0] in SYLLABLES else "garden"

    def logprob_steps(self, doc):
        """
        每个 token 的真实 log p(x_t | x_<t)（自然对数）。第一个 token 还要乘上文档类型的概率（0.6 或 0.4）。

        【思路】百科条目：名字的第 k 个音节的概率 = 以"前 k 个音节"开头的人物的总权重 / 以"前 k−1 个音节"开头的总权重；
                关系动词 1/3；之后（in、值、句号、<eos>）全部确定，概率 1。
                花园句子：沿第 1 讲的状态机逐词取概率。
        """
        if self.kind(doc) == "bio":
            nm = tuple(doc[:3])
            out = [math.log(P_BIO * self.prefix_mass[nm[:1]])]
            for k in (2, 3):
                out.append(math.log(self.prefix_mass[nm[:k]] / self.prefix_mass[nm[:k - 1]]))
            out.append(math.log(1 / 3))
            out += [0.0] * (len(doc) - 4)
            return out
        _, steps = self.garden.logprob(doc, return_steps=True)
        out = [math.log(p) for p in steps]
        out[0] += math.log(1 - P_BIO)
        return out


ALMANAC = Almanac()


def token_category(doc, t):
    """第 t 个 token 属于哪一类：'事实'（百科条目里的值）、'名字'（百科条目的 3 个音节）、'其他'。"""
    if Almanac.kind(doc) == "bio":
        if t < 3:
            return "名字"
        if doc[t] in TOWNS or doc[t] in PLANTS or doc[t] in ANIMALS:
            return "事实"
    return "其他"


def stream_batches(seed, batch_size=16, seq_len=64, unique_tokens=None):
    """
    无穷的训练批次：文档首尾相接（每篇前面加 <bos>），切成 seq_len + 1 长的片段（第 3 讲的"打包"）。

    【参数】seed：数据的随机种子（同一个种子 → 同样的数据顺序，所以不同大小的模型看到的是同一批数据）；
            unique_tokens：给定时只生成这么多 token 的语料，然后**反复**使用（04 号脚本的"数据受限"实验）；
                           每一遍的切片位置不同（错开一个随机偏移）。
    【返回】生成器，每次给出 (x, y)：形状都是 (batch_size, seq_len) 的 long 张量，y 是 x 错开一位。
    """
    rng = random.Random(seed)
    need = batch_size * (seq_len + 1)
    if unique_tokens is None:
        buf = []
        while True:
            while len(buf) < need:
                buf += [BOS] + VOCAB.encode(ALMANAC.sample_doc(rng))
            chunk, buf = buf[:need], buf[need:]
            t = torch.tensor(chunk).view(batch_size, seq_len + 1)
            yield t[:, :-1], t[:, 1:]
    data = []
    while len(data) < unique_tokens:
        data += [BOS] + VOCAB.encode(ALMANAC.sample_doc(rng))
    data = data[:unique_tokens]
    epoch = 0
    while True:
        off = (epoch * 7919) % (seq_len + 1)
        seqs = [data[i:i + seq_len + 1] for i in range(off, len(data) - seq_len, seq_len + 1)]
        random.Random(seed * 1000 + epoch).shuffle(seqs)
        for s in range(0, len(seqs) - batch_size + 1, batch_size):
            t = torch.tensor(seqs[s:s + batch_size])
            yield t[:, :-1], t[:, 1:]
        epoch += 1


def valid_docs(n=3000, seed=12345):
    """固定的验证集：n 篇新文档（逐篇评估，不打包，这样每个 token 的真实条件概率是精确的）。"""
    rng = random.Random(seed)
    return [ALMANAC.sample_doc(rng) for _ in range(n)]


VALID = valid_docs()


def true_entropy(docs=VALID):
    """
    验证集上"真实分布"的平均每 token 交叉熵（nats）：这是任何模型都不可能低于的下界，即不可约损失 E。

    【返回】dict：{"全部": E, "事实": …, "名字": …, "其他": …}。
    """
    tot = {"全部": [0.0, 0], "事实": [0.0, 0], "名字": [0.0, 0], "其他": [0.0, 0]}
    for d in docs:
        for t, lp in enumerate(ALMANAC.logprob_steps(d)):
            c = token_category(d, t)
            for k in ("全部", c):
                tot[k][0] -= lp
                tot[k][1] += 1
    return {k: s / n for k, (s, n) in tot.items()}


E_TRUE = true_entropy()

# =====================================================================
# 2. 模型族
# =====================================================================

# 名字 → (d_model, n_layers)。宽度 d，前馈层 4d，每头 16 维。非嵌入参数 ≈ 12·L·d²。
FAMILY = {"d8": (8, 2), "d12": (12, 2), "d16": (16, 2), "d24": (24, 2), "d32": (32, 2), "d48": (48, 2),
          "d64": (64, 2), "d96": (96, 2), "d128": (128, 2)}


# 每种宽度的峰值学习率：03 号脚本第 3 节的扫描（500k token，学习率 3e-3 / 1e-2 / 3e-2 / 6e-2）里最好的那个。
# 第一版所有大小都用 3e-3，小模型被严重拖累（d16：1.85 vs 调好后 1.58），规模定律的结论因此偏向大模型。
LR_BY_WIDTH = {8: 0.06, 12: 0.06, 16: 0.03, 24: 0.06, 32: 0.03, 48: 0.01, 64: 0.01, 96: 0.01, 128: 0.01}


def build(d, L, seed=0):
    from transformer import TransformerLM
    torch.manual_seed(seed)
    return TransformerLM(len(VOCAB), d_model=d, n_heads=max(1, d // 16), n_layers=L, d_ff=4 * d)


def count_params(model):
    """【返回】(总参数, 非嵌入参数)。输出层和输入嵌入共享权重，嵌入只算一次（Kaplan et al. 的 N 不含嵌入）。"""
    total = sum(p.numel() for p in model.parameters())
    return total, total - model.emb.weight.numel()


def n_params(d, L):
    return count_params(build(d, L))


# =====================================================================
# 3. 训练（带结果缓存）
# =====================================================================

def lr_at(step, total, warmup, peak, floor=0.1):
    """预热 + 余弦衰减到 floor × 峰值（Chinchilla 的做法：余弦周期和训练长度一致）。"""
    if step < warmup:
        return peak * (step + 1) / warmup
    prog = min(1.0, (step - warmup) / max(1, total - warmup))
    return peak * (floor + (1 - floor) * 0.5 * (1 + math.cos(math.pi * prog)))


@torch.no_grad()
def evaluate(model, docs=VALID, by_category=False):
    """
    验证集上的平均每 token 交叉熵（nats）。逐篇评估（按长度分组成批，不需要补齐）。

    【返回】float；by_category=True 时返回 dict（同 true_entropy 的键）。
    """
    model.eval()
    tot = {"全部": [0.0, 0], "事实": [0.0, 0], "名字": [0.0, 0], "其他": [0.0, 0]}
    groups = {}
    for d in docs:
        groups.setdefault(len(d), []).append(d)
    for L in sorted(groups):
        g = groups[L]
        for s in range(0, len(g), 512):
            part = g[s:s + 512]
            ids = torch.tensor([[BOS] + VOCAB.encode(d) for d in part])
            nll = F.cross_entropy(model(ids[:, :-1]).transpose(1, 2), ids[:, 1:], reduction="none")
            for row, d in zip(nll.tolist(), part):
                for t, x in enumerate(row):
                    c = token_category(d, t)
                    for k in ("全部", c):
                        tot[k][0] += x
                        tot[k][1] += 1
    out = {k: s / n for k, (s, n) in tot.items()}
    return out if by_category else out["全部"]


def _key(cfg):
    return hashlib.sha1(json.dumps(cfg, sort_keys=True).encode()).hexdigest()[:12]


def _load_results():
    if os.path.exists(RESULTS):
        with open(RESULTS) as f:
            return json.load(f)
    return {}


def run(d, L, tokens, lr=None, batch_size=16, seq_len=64, seed=0, schedule_tokens=None, unique_tokens=None,
        eval_points=(), save=False, verbose=True):
    """
    训练一个模型并返回结果（有缓存就直接读 results.json，不重新训练）。

    【参数】d, L：宽度、层数；tokens：训练 token 数 D；lr：峰值学习率（默认按宽度查 LR_BY_WIDTH）；
            schedule_tokens：余弦周期按多少 token 设定（默认 = tokens，即周期和训练长度一致；
                             Kaplan et al. 用的是固定的长周期，03 号脚本对比两者）；
            unique_tokens：数据受限时不重复的 token 数（默认不限）；
            eval_points：在这些 token 数处额外评估一次验证损失（画训练曲线用）；
            save：是否把训练好的参数存进 checkpoints/（05 号脚本要用模型本身）。
    【返回】dict：N（非嵌入参数）、N_total、D、C（= 6·N·D）、loss（最终验证损失）、loss_by（按类别）、
            curve（[(token 数, 验证损失)]）、seconds。
    【边界】缓存的键 = 全部参数；改了训练代码后请删除 results.json 和 checkpoints/。
    """
    lr = LR_BY_WIDTH[d] if lr is None else lr
    cfg = {"d": d, "L": L, "tokens": tokens, "lr": lr, "B": batch_size, "T": seq_len, "seed": seed,
           "sched": schedule_tokens, "unique": unique_tokens, "evals": list(eval_points), "v": 1}
    key = _key(cfg)
    res = _load_results()
    path = os.path.join(CKPT_DIR, f"{key}.pt")
    if key in res and (not save or os.path.exists(path)):
        return res[key]
    model = build(d, L, seed)
    n_tot, n_ne = count_params(model)
    steps = max(1, tokens // (batch_size * seq_len))
    sched_steps = max(1, (schedule_tokens or tokens) // (batch_size * seq_len))
    warmup = max(1, min(100, sched_steps // 10))
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    data = stream_batches(seed, batch_size, seq_len, unique_tokens)
    eval_steps = {max(1, e // (batch_size * seq_len)): e for e in eval_points}
    curve = []
    t0 = time.time()
    model.train()
    for step in range(steps):
        for g in opt.param_groups:
            g["lr"] = lr_at(step, sched_steps, warmup, lr)
        x, y = next(data)
        loss = F.cross_entropy(model(x).reshape(-1, len(VOCAB)), y.reshape(-1))
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        if step + 1 in eval_steps:
            curve.append((eval_steps[step + 1], evaluate(model)))
            model.train()
    by = evaluate(model, by_category=True)
    out = {"d": d, "L": L, "lr": lr, "N": n_ne, "N_total": n_tot, "D": steps * batch_size * seq_len,
           "C": 6.0 * n_ne * steps * batch_size * seq_len, "loss": by["全部"], "loss_by": by, "curve": curve,
           "seconds": round(time.time() - t0, 1)}
    if verbose:
        print(f"    训练 d={d} L={L}（N={n_ne:,}）D={out['D']:,}：验证损失 {out['loss']:.4f}，用时 {out['seconds']} s")
    if save:
        os.makedirs(CKPT_DIR, exist_ok=True)
        torch.save(model.state_dict(), path)
    res = _load_results()
    res[key] = out
    with open(RESULTS, "w") as f:
        json.dump(res, f, ensure_ascii=False, indent=0, sort_keys=True)
    return out


def load_model(d, L, tokens, **kw):
    """读回 run(..., save=True) 存下的模型（05 号脚本用）。"""
    cfg = {"d": d, "L": L, "tokens": tokens, "lr": kw.get("lr") or LR_BY_WIDTH[d], "B": kw.get("batch_size", 16),
           "T": kw.get("seq_len", 64), "seed": kw.get("seed", 0), "sched": kw.get("schedule_tokens"),
           "unique": kw.get("unique_tokens"), "evals": list(kw.get("eval_points", ())), "v": 1}
    run(d, L, tokens, save=True, verbose=kw.get("verbose", True),
        **{k: v for k, v in kw.items() if k != "verbose"})
    m = build(d, L)
    m.load_state_dict(torch.load(os.path.join(CKPT_DIR, f"{_key(cfg)}.pt")))
    m.eval()
    return m


# =====================================================================
# 4. 拟合
# =====================================================================

def fit_power(x, y, E=0.0):
    """
    拟合 y − E = A · x^(−α)：在对数坐标上做最小二乘直线。

    【参数】x, y：数组；E：已知（或假设）的不可约部分。
    【返回】(A, α)。
    【思路】log(y − E) = log A − α log x 是一条直线；斜率的相反数就是幂律指数。
    """
    lx, ly = np.log(np.asarray(x, float)), np.log(np.asarray(y, float) - E)
    slope, icpt = np.polyfit(lx, ly, 1)
    return float(np.exp(icpt)), float(-slope)


def fit_power_offset(x, y):
    """
    拟合 y = E + A · x^(−α)，E 也是未知数：在 E 的网格上逐个做 fit_power，取残差最小的。

    【返回】(E, A, α)。
    【边界】E 只在 [0, min(y)) 里搜索；数据点太少或曲线太直时 E 很难估准（02 号脚本会看到）。
    """
    y = np.asarray(y, float)
    best = None
    for E in np.linspace(0, y.min() * 0.999, 400):
        A, a = fit_power(x, y, E)
        pred = E + A * np.asarray(x, float) ** (-a)
        err = float(np.sum((np.log(pred) - np.log(y)) ** 2))
        if best is None or err < best[0]:
            best = (err, E, A, a)
    return best[1], best[2], best[3]


def fit_chinchilla(N, D, L, delta=1e-3):
    """
    Chinchilla 的方法 3：拟合 L(N, D) = E + A / N^α + B / D^β。

    【思路】和原论文一样：参数化成 E = e^e、A = e^a、B = e^b，在 log 空间里用 Huber 损失
            Huber_δ( log L̂ − log L ) 拟合（对离群点稳健），从一组初始值出发各跑一次 L-BFGS，取最好的。
    【返回】dict：E、A、alpha、B、beta。
    """
    N, D, L = (torch.tensor(np.asarray(v, float)) for v in (N, D, L))
    lN, lD, lL = N.log(), D.log(), L.log()
    best = None
    for a0 in (0.0, 5.0, 10.0):
        for b0 in (0.0, 5.0, 10.0):
            for e0 in (-1.0, 0.0, 0.5):
                for al0 in (0.2, 0.5):
                    p = torch.tensor([e0, a0, b0, al0, al0], dtype=torch.float64, requires_grad=True)
                    opt = torch.optim.LBFGS([p], max_iter=200, line_search_fn="strong_wolfe")

                    def closure():
                        opt.zero_grad()
                        pred = torch.logsumexp(torch.stack([p[0].expand_as(lN), p[1] - p[3] * lN,
                                                            p[2] - p[4] * lD]), 0)
                        loss = F.huber_loss(pred, lL, delta=delta, reduction="sum")
                        loss.backward()
                        return loss

                    try:
                        opt.step(closure)
                    except RuntimeError:
                        continue
                    with torch.no_grad():
                        pred = torch.logsumexp(torch.stack([p[0].expand_as(lN), p[1] - p[3] * lN,
                                                            p[2] - p[4] * lD]), 0)
                        loss = float(F.huber_loss(pred, lL, delta=delta, reduction="sum"))
                    if math.isfinite(loss) and (best is None or loss < best[0]):
                        best = (loss, p.detach().clone())
    e, a, b, al, be = best[1].tolist()
    return {"E": math.exp(e), "A": math.exp(a), "alpha": al, "B": math.exp(b), "beta": be}


def chinchilla_loss(f, N, D):
    return f["E"] + f["A"] / N ** f["alpha"] + f["B"] / D ** f["beta"]


def chinchilla_optimum(f, C):
    """
    给定拟合结果和算力 C（= 6ND），使 L 最小的 N*、D*（解析解，Hoffmann et al. 2022 附录 D）：
        N* = G · (C/6)^(β/(α+β)),  D* = (1/G) · (C/6)^(α/(α+β)),  G = (αA / (βB))^(1/(α+β))。
    """
    a, b = f["alpha"], f["beta"]
    G = (a * f["A"] / (b * f["B"])) ** (1 / (a + b))
    return G * (C / 6) ** (b / (a + b)), (1 / G) * (C / 6) ** (a / (a + b))


__all__ = ["ALMANAC", "VOCAB", "VALID", "E_TRUE", "FAMILY", "SYLLABLES", "TOWNS", "PLANTS", "ANIMALS",
           "build", "count_params", "n_params", "run", "load_model", "evaluate", "stream_batches",
           "fit_power", "fit_power_offset", "fit_chinchilla", "chinchilla_loss", "chinchilla_optimum",
           "true_entropy", "token_category", "N_PEOPLE", "lr_at", "LR_BY_WIDTH"]