"""
scaling_lab.py —— 规模定律的小实验室：能精确算出"不可约损失"的 Zipf 事实世界、小 GPT 的训练、幂律拟合（第 15 讲）
=====================================================================

【这个文件要回答的问题】
    slides：大语言模型的规模定律通常假设损失是一个幂律 f(x) = a·x^{−k}；Kaplan et al. (2020) 和 Hoffmann et al. (2022,
    Chinchilla) 都在拟合"测试损失怎样随参数量 N、训练 token 数 D、计算量 C 变化"。真实的语料里，**不可约损失**
    （语言本身的熵）是未知的，只能和别的参数一起拟合出来。
    这里造一个**能精确算出不可约损失**的小世界：
        K = 4096 个实体，每个实体的名字是 3 个"音节"（16 个音节里选，16³ = 4096，名字和实体一一对应）；
        每个实体有 A = 4 个属性，每个属性的值是 V = 16 个里随机定好的一个（"事实"）；
        一句话 = <s> 音节₁ 音节₂ 音节₃ 属性 值 </s>，实体按 Zipf 分布出现（p_e ∝ 排名^{−a}），属性均匀。
    于是"真实分布"完全已知：语言的熵 H = H(实体) + log A（值和 </s> 都是确定的）。模型分布和真实分布的 KL 散度
    （超出不可约损失的那一部分）可以对全部 K × A = 16384 种句子**精确**加权求和，没有抽样噪声。
    Zipf 分布让"记住更多事实"永远有用（越往后的事实越少见）——这正是幂律出现的原因（03 号脚本的理论）。

【本文件提供】
    FactWorld            世界：sample / packed（训练用的拼接序列）/ all_sentences（精确评估用）/ entropy
    make_model / param_counts      第 2 讲的 GPT；参数量分"全部"和"非嵌入"（Kaplan 只数非嵌入的参数）
    train_lm             余弦学习率（长度和 D 匹配，Chinchilla 强调过这一点）；数据来源可以换（05 号脚本的"过滤 vs 不过滤"）
    exact_kl             精确的超额损失（nats / 句），并拆成名字、属性、值、</s> 四部分
    fit_power / fit_chinchilla / optimal_allocation      幂律拟合、L(N, D) = E + A/N^α + B/D^β 的拟合、计算最优配比
    cached_run           训练结果（不是权重）缓存成 results/*.json，重复运行直接读取
"""

import json
import math
import os
import sys
import time

import torch
import torch.nn.functional as F

_HERE = os.path.dirname(os.path.abspath(__file__))
_p = os.path.join(_HERE, "..", "lecture02")
if _p not in sys.path:
    sys.path.insert(0, _p)
from lm_models import GPT  # noqa: E402

RESULTS_DIR = os.path.join(_HERE, "results")
T_SEQ = 63                                  # 训练序列长度：9 句话 × 7 个 token


# =====================================================================
# 1. 世界
# =====================================================================

class FactWorld:
    """
    【参数】K：实体数（必须 ≤ 16³ = 4096）；A：每个实体的属性数；V：每个属性的取值数；a：Zipf 指数；seed：名字和事实的随机种子
    【词表】0 = <s>，1 = </s>，2..17 = 16 个音节，接着 A 个属性词，最后 A × V 个值（每个属性的值各自独立编号）
    【属性】probs (K,) float64：实体的出现概率（按排名，排名 0 最常见）；names (K, 3)；facts (K, A)
    """

    BOS, EOS, SYL0, N_SYL = 0, 1, 2, 16

    def __init__(self, K=4096, A=4, V=16, a=1.2, seed=0):
        if K > self.N_SYL ** 3:
            raise ValueError("K 不能超过 16³ = 4096")
        self.K, self.A, self.V, self.a = K, A, V, a
        g = torch.Generator().manual_seed(seed)
        p = torch.arange(1, K + 1, dtype=torch.float64) ** (-a)
        self.probs = p / p.sum()
        self.cum = self.probs.cumsum(0)
        code = torch.randperm(self.N_SYL ** 3, generator=g)[:K]          # 实体排名 → 名字（随机打乱，名字和频率无关）
        self.names = torch.stack([code // 256, (code // 16) % 16, code % 16], 1)
        self.facts = torch.randint(0, V, (K, A), generator=g)
        self.ATT0 = self.SYL0 + self.N_SYL
        self.VAL0 = self.ATT0 + A
        self.vocab_size = self.VAL0 + A * V

    def entropy(self):
        """语言的熵（nats / 句）= H(实体) + log A：模型能达到的最低损失（不可约损失）。"""
        p = self.probs
        return float(-(p * p.log()).sum() + math.log(self.A))

    def sentences(self, e, at, values=None):
        """实体编号 e、属性编号 at（都是 (n,) long）→ (n, 7) 的 token。values 不给就用真实的事实。"""
        v = self.facts[e, at] if values is None else values
        return torch.stack([torch.full_like(at, self.BOS), self.SYL0 + self.names[e, 0], self.SYL0 + self.names[e, 1],
                            self.SYL0 + self.names[e, 2], self.ATT0 + at, self.VAL0 + at * self.V + v,
                            torch.full_like(at, self.EOS)], 1)

    def sample(self, n, gen, noise=0.0):
        """
        按真实分布抽 n 句话。noise > 0 时，每句话以这个概率把值换成一个随机的值（"低质量数据"：错的事实）。
        【返回】(n, 7) long
        """
        e = torch.searchsorted(self.cum, torch.rand(n, generator=gen, dtype=torch.float64)).clamp(max=self.K - 1)
        at = torch.randint(0, self.A, (n,), generator=gen)
        values = None
        if noise:
            wrong = torch.rand(n, generator=gen) < noise
            values = torch.where(wrong, torch.randint(0, self.V, (n,), generator=gen), self.facts[e, at])
        return self.sentences(e, at, values)

    def all_sentences(self):
        """全部 K × A 种（正确的）句子和它们的真实概率 p_e / A。【返回】((K·A, 7) long, (K·A,) float64)"""
        e = torch.arange(self.K).repeat_interleave(self.A)
        at = torch.arange(self.A).repeat(self.K)
        return self.sentences(e, at), self.probs[e] / self.A


def packed(sents, B):
    """把 B × 9 句话拼成 B 条长 63 的训练序列：输入是前 63 个 token，目标是后 63 个（每条序列以 <s> 开头）。"""
    s = sents[:B * 9].reshape(B, 63)
    s = torch.cat([s, torch.zeros(B, 1, dtype=torch.long)], 1)          # 补一个 <s>：最后一个 </s> 之后预测下一句的开头
    return s[:, :-1], s[:, 1:]


# =====================================================================
# 2. 模型与训练
# =====================================================================

def make_model(world, d, n_layers=2):
    """第 2 讲的 GPT：d 维、n_layers 层、d ≥ 32 时 4 个头（否则 2 个），前馈层 4d。构造前固定种子。"""
    torch.manual_seed(0)
    return GPT(world.vocab_size, d, n_layers, 4 if d >= 32 else 2, max_len=64)


def param_counts(model):
    """
    【返回】(全部参数, 非嵌入参数)。"嵌入" = 词嵌入表 + 位置嵌入表 + 输出层（都和词表或位置的个数成正比）。
    Kaplan et al. 只数非嵌入参数；Chinchilla 数全部参数——这是两者结论不同的原因之一（Pearce & Song, 2024；04 号脚本）。
    """
    total = sum(p.numel() for p in model.parameters())
    emb = sum(p.numel() for n, p in model.named_parameters() if n.startswith(("tok", "pos", "head")))
    return total, total - emb


def train_lm(model, world, tokens, batch=64, lr=3e-3, warmup=100, seed=0, source=None):
    """
    在世界里训练语言模型 tokens 个 token（每步 batch × 63 个）。
    【学习率】先线性预热 min(warmup, 总步数 / 10) 步，再按余弦从 lr 降到 0.1·lr——余弦的长度等于这次训练的总步数。
        Hoffmann et al. 指出：Kaplan 用的余弦长度固定、不随 D 变，小 D 的实验没降完、大 D 的又早早降完，这会扭曲规模曲线。
    【参数】source：None 表示每一步都从真实分布抽新的句子（"数据无限、每个 token 只见一次"）；
        也可以是一个函数 source(gen, n) → (n, 7)，05 号脚本用它换成"过滤后的有限数据"或"带噪声的数据"
    【返回】{"steps", "seconds"}
    """
    steps = max(1, round(tokens / (batch * T_SEQ)))
    warmup = max(1, min(warmup, steps // 10))
    torch.manual_seed(seed)
    gen = torch.Generator().manual_seed(seed + 1)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.0)
    model.train()
    t0 = time.perf_counter()
    for st in range(steps):
        frac = st / max(steps - 1, 1)
        for gr in opt.param_groups:
            gr["lr"] = lr * min(1.0, (st + 1) / warmup) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * frac)))
        sents = world.sample(batch * 9, gen) if source is None else source(gen, batch * 9)
        x, y = packed(sents, batch)
        logits, _ = model(x)
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
    model.eval()
    return {"steps": steps, "seconds": time.perf_counter() - t0}


@torch.no_grad()
def exact_kl(model, world, batch=4096):
    """
    精确的超额损失：KL(真实分布 ‖ 模型) = E_真实[−log q(句子)] − H（nats / 句），对全部 K × A 种句子按真实概率加权。
    【返回】{"total", "names", "attr", "value", "eos"}：
        names  E[−log q(三个音节)] − H(实体)：模型把实体的频率学得多准
        attr   E[−log q(属性 | 名字)] − log A
        value  E[−log q(值 | 名字, 属性)]：事实记住了多少（真实分布里值是确定的，所以这一项本身就是超额损失）
        eos    E[−log q(</s> | …)]
    【边界】评估时每句话单独从 <s> 开始（训练时一句话前面可能还有别的句子）；训练序列总是以一句完整的话开头，
        所以模型见过这种情形。
    """
    sents, w = world.all_sentences()
    parts = torch.zeros(6, dtype=torch.float64)
    for i in range(0, len(sents), batch):
        s = sents[i:i + batch]
        logp = F.log_softmax(model(s[:, :-1])[0].double(), -1).gather(2, s[:, 1:, None])[..., 0]   # (n, 6)
        parts += (-logp * w[i:i + batch, None]).sum(0)
    h_ent = float(-(world.probs * world.probs.log()).sum())
    names = float(parts[:3].sum()) - h_ent
    attr = float(parts[3]) - math.log(world.A)
    value, eos = float(parts[4]), float(parts[5])
    return {"total": names + attr + value + eos, "names": names, "attr": attr, "value": value, "eos": eos}


def cached_run(key, fn, log=print):
    """
    训练结果（一个只含数字的字典）缓存在 results/<key>.json：有就直接读取，没有就 fn() 算出来再存。
    【为什么只存结果不存权重】规模定律的实验要训练很多个小模型，我们只关心它们的损失，不需要再用这些模型。
    """
    path = os.path.join(RESULTS_DIR, f"{key}.json")
    if os.path.exists(path):
        with open(path) as f:
            return json.load(f)
    out = fn()
    os.makedirs(RESULTS_DIR, exist_ok=True)
    with open(path, "w") as f:
        json.dump(out, f, indent=1, sort_keys=True)
    if log:
        log(f"    （已保存 results/{key}.json）")
    return out


def run_point(world, d, tokens, key_prefix="clean", source=None, n_layers=2, log=print, cache=True):
    """
    训练一个 (宽度 d, 数据量 tokens) 的模型并精确评估。
    【参数】key_prefix：缓存文件名的前缀（不同的数据来源要用不同的前缀）；cache=False：不读写缓存（冒烟测试用）
    【返回】{"d", "tokens", "N", "N_nonemb", "C", **exact_kl, "steps", "seconds"}
    """
    key = f"{key_prefix}_K{world.K}_a{world.a}_d{d}_L{n_layers}_D{int(tokens)}"

    def fn():
        m = make_model(world, d, n_layers)
        info = train_lm(m, world, tokens, source=source)
        N, N_ne = param_counts(m)
        D = info["steps"] * 64 * T_SEQ
        return dict(d=d, tokens=D, N=N, N_nonemb=N_ne, C=6 * N * D, steps=info["steps"],
                    seconds=round(info["seconds"], 1), **exact_kl(m, world))
    return cached_run(key, fn, log=log) if cache else fn()


# =====================================================================
# 3. 拟合
# =====================================================================

def fit_power(x, y):
    """
    拟合 y = A·x^{−α}：两边取对数，log y = log A − α log x 是一条直线，用最小二乘求斜率。
    【返回】(A, α, R²)
    """
    lx = torch.tensor([math.log(v) for v in x], dtype=torch.float64)
    ly = torch.tensor([math.log(v) for v in y], dtype=torch.float64)
    X = torch.stack([torch.ones_like(lx), lx], 1)
    coef = torch.linalg.lstsq(X, ly[:, None]).solution[:, 0]
    pred = X @ coef
    r2 = 1 - float(((ly - pred) ** 2).sum() / ((ly - ly.mean()) ** 2).sum())
    return math.exp(float(coef[0])), -float(coef[1]), r2


def fit_chinchilla(N, D, L, E_fixed=None, delta=1e-3):
    """
    Chinchilla 的"方法 3"：拟合 L(N, D) = E + A/N^α + B/D^β。和论文一样，在 log 空间用 Huber 损失、
    从一组初始值出发各跑一次 L-BFGS，取最好的那个。
    【参数】N、D、L：同样长度的列表；E_fixed：给定不可约损失（我们的世界里它已知），None 表示一起拟合
    【返回】{"E", "A", "B", "alpha", "beta", "loss"}
    【边界】点数少时（比如 12 个点拟 5 个参数）拟合很不稳定，04 号脚本会把"固定 E"和"一起拟合 E"都列出来。
    """
    lN = torch.tensor([math.log(v) for v in N], dtype=torch.float64)
    lD = torch.tensor([math.log(v) for v in D], dtype=torch.float64)
    lL = torch.tensor([math.log(v) for v in L], dtype=torch.float64)
    best = None
    for a0 in (0.0, 0.5, 1.0):
        for b0 in (0.0, 0.5, 1.0):
            for e0 in ((-1.0, 0.5) if E_fixed is None else (None,)):
                th = torch.tensor([5.0, 5.0, a0, b0] + ([] if e0 is None else [e0]), dtype=torch.float64,
                                  requires_grad=True)
                opt = torch.optim.LBFGS([th], max_iter=500, line_search_fn="strong_wolfe")

                def closure():
                    opt.zero_grad()
                    a_, b_, alpha, beta = th[0], th[1], th[2], th[3]
                    lE = th[4] if E_fixed is None else torch.tensor(math.log(E_fixed), dtype=torch.float64)
                    pred = torch.logsumexp(torch.stack([a_ - alpha * lN, b_ - beta * lD, lE.expand_as(lN)]), 0)
                    r = pred - lL
                    loss = torch.where(r.abs() < delta, 0.5 * r ** 2, delta * (r.abs() - 0.5 * delta)).sum()
                    loss.backward()
                    return loss

                try:
                    loss = float(opt.step(closure).detach())
                except RuntimeError:
                    continue
                if math.isfinite(loss) and (best is None or loss < best[0]):
                    best = (loss, th.detach().clone())
    loss, th = best
    E = E_fixed if E_fixed is not None else math.exp(float(th[4]))
    return {"E": E, "A": math.exp(float(th[0])), "B": math.exp(float(th[1])), "alpha": float(th[2]),
            "beta": float(th[3]), "loss": loss}


def optimal_allocation(fit, C):
    """
    给定拟合 L = E + A/N^α + B/D^β 和计算量 C = 6ND，损失最小的 N、D（Chinchilla 的公式 (4)）：
        N_opt = G·(C/6)^{β/(α+β)}，D_opt = G^{−1}·(C/6)^{α/(α+β)}，G = (αA / (βB))^{1/(α+β)}
    【返回】(N_opt, D_opt)
    """
    a, b = fit["alpha"], fit["beta"]
    G = (a * fit["A"] / (b * fit["B"])) ** (1 / (a + b))
    return G * (C / 6) ** (b / (a + b)), (C / 6) ** (a / (a + b)) / G


__all__ = ["FactWorld", "packed", "make_model", "param_counts", "train_lm", "exact_kl", "cached_run", "run_point",
           "fit_power", "fit_chinchilla", "optimal_allocation", "RESULTS_DIR", "T_SEQ"]