"""
=============================================================================
CMU 17-630 · 第 1 讲共享模块：tiny_gpt.py —— 一个能在笔记本 CPU 上训练的迷你 GPT
=============================================================================

【这个模块是做什么的】
    GPT 系列（GPT-1/2/3……）的结构其实很简单：一叠"只看左边"的 Transformer 层，
    训练目标只有一个——**预测下一个 token**。这个模块用 PyTorch 从零写一个同样结构、
    但小几十万倍的模型（约 26 万个参数），在 pretrain_world.py 的迷你语料上训练不到一分钟，
    然后我们就可以像对待真正的 base 模型一样去"提示"它：

        model, tok = get_model("base")
        generate(model, tok, "<doc> faq : q : what is the capital of varnia ? a :")

    它和真的 GPT 只差在规模上：层数（2 vs 96）、宽度（96 vs 12288）、语料（23 万词 vs 3000 亿 token）。
    所以它能展示"续写文档"这件事的**机制**——格式敏感、知识跨格式迁移、温度采样——
    但不能展示只有大模型才有的能力（例如复杂的少样本学习）。这一点讲义里会反复提醒。

【模块里有什么】
    GPTConfig                   模型的超参数
    CausalSelfAttention, Block  Transformer 层（手写注意力，方便后面几讲把注意力权重拿出来看）
    TinyGPT                     整个模型：词嵌入 + 位置嵌入 + N 层 + 输出层
    train_model                 训练循环（AdamW、学习率预热 + 余弦衰减）
    get_model                   按名字取一个标准模型：有缓存就读缓存，没有就训练并存下来
    next_token_probs / top_next 下一个 token 的分布
    score_options               几个候选词谁的概率大（用来做分类）
    generate                    续写（贪心或按温度采样）
    hidden_states               最后一层每个位置的向量（上下文相关的表示）
    perplexity                  困惑度

【可复现性】
    torch.set_num_threads(1) + 固定随机种子 + 固定的数据顺序：同一台机器上重新训练，
    得到的权重逐位相同（换一台机器、换 PyTorch 版本时可能有极小的浮点差异）。

【运行】python3 tiny_gpt.py            # 训练（或读取缓存）标准模型 "base"，打印几个续写
        python3 tiny_gpt.py --quick    # 只训练 60 步的 smoke test，几秒钟，不写缓存
【依赖】PyTorch。第 1 讲的 05、06 号脚本，以及后面几讲（例如第 3 讲的指令微调）都会用它。
"""

import hashlib
import json
import math
import os
import sys
import time
from dataclasses import asdict, dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

from pretrain_world import WordTokenizer, build_corpus, build_world, corpus_text

HERE = os.path.dirname(os.path.abspath(__file__))
CKPT_DIR = os.path.join(HERE, "checkpoints")


# ---------------------------------------------------------------------------
# 1. 模型
# ---------------------------------------------------------------------------

@dataclass
class GPTConfig:
    """
    模型超参数。

    【字段】vocab_size 词表大小；block_size 上下文窗口（最多看前面多少个 token）；
        n_layer 层数；n_head 注意力头数；d_model 每个 token 的向量维度；
        mlp_ratio MLP 隐藏层是 d_model 的几倍。
    """
    vocab_size: int
    block_size: int = 64
    n_layer: int = 2
    n_head: int = 4
    d_model: int = 96
    mlp_ratio: int = 4


class CausalSelfAttention(nn.Module):
    """
    因果自注意力（causal self-attention）：每个位置只能看见它自己和它左边的位置。

    【思路】对每个位置 t：
        q_t = x_t W_q（"我在找什么"），k_s = x_s W_k（"我这里有什么"），v_s = x_s W_v（"我能提供什么"）
        权重 a_{t,s} = softmax_s( q_t · k_s / sqrt(d_head) )，只对 s ≤ t 计算（右边的位置被遮住）
        输出 = Σ_s a_{t,s} v_s
        多头（multi-head）= 把向量切成 n_head 段，每段各算一份注意力，再拼起来——
        不同的头可以学会"看"不同的东西（例如一个头看前一个词，另一个头找同样的词）。
    【为什么手写而不用 F.scaled_dot_product_attention】为了能把注意力权重 a 返回出来看
        （第 2 讲会用它找"归纳头"）。
    """

    def __init__(self, cfg):
        super().__init__()
        assert cfg.d_model % cfg.n_head == 0
        self.n_head, self.d_head = cfg.n_head, cfg.d_model // cfg.n_head
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model)    # 一次算出 q、k、v
        self.proj = nn.Linear(cfg.d_model, cfg.d_model)       # 多头拼接后再混合一次
        # 下三角矩阵：mask[t, s] = 1 表示位置 t 可以看位置 s（s ≤ t）
        # 【Python 小知识】register_buffer 把一个张量登记为模型的一部分（会随模型保存、移动），
        #   但它不是可训练参数（不会被优化器更新）。
        self.register_buffer("mask", torch.tril(torch.ones(cfg.block_size, cfg.block_size)).bool())

    def forward(self, x, return_attn=False):
        """
        【参数】x：形状 (B, T, d_model) 的张量（B 个序列，每个 T 个位置）
        【返回】(输出 (B, T, d_model), 注意力权重 (B, n_head, T, T) 或 None)
        【Python 小知识】.view(B, T, n_head, d_head).transpose(1, 2) 把最后一维切成多头，
            再把"头"这一维挪到前面，得到 (B, n_head, T, d_head)。
        """
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        q = q.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)          # (B, h, T, T)
        att = att.masked_fill(~self.mask[:T, :T], float("-inf"))          # 遮住右边
        att = F.softmax(att, dim=-1)
        y = (att @ v).transpose(1, 2).contiguous().view(B, T, C)          # 拼回 (B, T, C)
        return self.proj(y), (att if return_attn else None)


class Block(nn.Module):
    """
    一个 Transformer 层："先注意力、再 MLP"，每一步前面有 LayerNorm，后面有残差连接。

    【思路】x = x + Attn(LN(x))：让每个位置从前文里"取"信息；
            x = x + MLP(LN(x))：在每个位置上独立地"加工"信息（大量知识被认为存在 MLP 里）。
        残差连接（x + ...）让信息和梯度可以直接穿过很多层；这是 GPT-2 起通用的 "pre-LN" 写法。
    """

    def __init__(self, cfg):
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.d_model)
        self.attn = CausalSelfAttention(cfg)
        self.ln2 = nn.LayerNorm(cfg.d_model)
        self.mlp = nn.Sequential(
            nn.Linear(cfg.d_model, cfg.mlp_ratio * cfg.d_model),
            nn.GELU(),
            nn.Linear(cfg.mlp_ratio * cfg.d_model, cfg.d_model),
        )

    def forward(self, x, return_attn=False):
        a, att = self.attn(self.ln1(x), return_attn)
        x = x + a
        x = x + self.mlp(self.ln2(x))
        return x, att


class TinyGPT(nn.Module):
    """
    迷你 GPT：词嵌入 + 位置嵌入 → n_layer 个 Block → LayerNorm → 输出层（和词嵌入共享权重）。

    【为什么输出层和词嵌入共享权重（weight tying）】输出层要给每个词打分，词嵌入要给每个词一个向量，
        两者都是"每个词一行"的矩阵；共享可以少一半参数，小模型上通常还更好（Press & Wolf 2017）。
    """

    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.pos_emb = nn.Embedding(cfg.block_size, cfg.d_model)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.ln_f = nn.LayerNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.head.weight = self.tok_emb.weight          # 共享权重
        self.apply(self._init_weights)

    @staticmethod
    def _init_weights(m):
        """按 GPT-2 的习惯初始化：线性层和嵌入用 N(0, 0.02²)，偏置为 0。"""
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.zeros_(m.bias)

    def n_params(self):
        """参数个数（共享的权重只算一次）。"""
        return sum(p.numel() for p in self.parameters())

    def features(self, idx, return_attn=False):
        """
        最后一层（LayerNorm 之后）每个位置的向量——模型对"到这里为止的全部前文"的理解。

        【返回】(形状 (B, T, d_model) 的张量, 每层的注意力列表)
        【提示工程联系】输出层就是拿这个向量和每个词的嵌入做点积来打分的。03 号脚本把它当作
            "上下文相关的句子向量"来用。
        """
        B, T = idx.shape
        assert T <= self.cfg.block_size, f"序列长度 {T} 超过了上下文窗口 {self.cfg.block_size}"
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        atts = []
        for blk in self.blocks:
            x, att = blk(x, return_attn)
            atts.append(att)
        return self.ln_f(x), atts

    def forward(self, idx, targets=None, return_attn=False):
        """
        【参数】idx：(B, T) 的 token id；targets：(B, T) 的"下一个 token"，给了就计算交叉熵损失
        【返回】(logits (B, T, V), loss 或 None, 每层的注意力列表或 None)
        【边界】T 不能超过 block_size（位置嵌入只有 block_size 个）。调用方负责截断。
        """
        h, atts = self.features(idx, return_attn)
        logits = self.head(h)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss, (atts if return_attn else None)


# ---------------------------------------------------------------------------
# 2. 训练
# ---------------------------------------------------------------------------

def get_batch(data, block_size, batch_size, gen):
    """
    从一条长 token 流里随机截取 batch_size 段长度为 block_size 的窗口。

    【返回】(x, y)：y 是 x 向右错开一位——"每个位置的答案就是下一个 token"
    【Python 小知识】torch.randint(high, (n,), generator=gen) 用指定的随机数生成器产生 n 个
        [0, high) 的整数；用独立的 generator 而不是全局随机数，结果才可复现。
    """
    ix = torch.randint(len(data) - block_size - 1, (batch_size,), generator=gen)
    x = torch.stack([data[i:i + block_size] for i in ix.tolist()])
    y = torch.stack([data[i + 1:i + 1 + block_size] for i in ix.tolist()])
    return x, y


def lr_at(step, total, base_lr, warmup):
    """学习率：前 warmup 步线性升到 base_lr，之后按余弦曲线降到 base_lr 的 10%。"""
    if step < warmup:
        return base_lr * (step + 1) / warmup
    t = (step - warmup) / max(1, total - warmup)
    return base_lr * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * t)))


def train_model(model, train_ids, steps=3000, batch_size=32, lr=3e-3, warmup=100,
                seed=0, eval_ids=None, log_every=500, verbose=True):
    """
    训练循环：反复"取一批 → 前向 → 算损失 → 反向传播 → 更新参数"。

    【参数】model：TinyGPT；train_ids / eval_ids：一维 LongTensor；steps：更新步数；
        batch_size：每批窗口数；lr：峰值学习率；warmup：预热步数；seed：数据顺序的种子
    【返回】历史记录列表 [{"step", "train_loss", "eval_loss"}, ...]
    【边界】eval_ids 为 None 时不算验证损失。梯度范数裁剪到 1.0，防止偶尔的大梯度把训练搞崩。
    【Python 小知识】with torch.no_grad(): 里面的计算不记录梯度——评估时不需要梯度，省内存也更快。
    """
    gen = torch.Generator().manual_seed(seed)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.1)
    hist, t0 = [], time.time()
    for step in range(steps):
        for g in opt.param_groups:
            g["lr"] = lr_at(step, steps, lr, warmup)
        model.train()
        x, y = get_batch(train_ids, model.cfg.block_size, batch_size, gen)
        _, loss, _ = model(x, y)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        if (step + 1) % log_every == 0 or step == steps - 1:
            rec = {"step": step + 1, "train_loss": loss.item()}
            if eval_ids is not None:
                rec["eval_loss"] = math.log(perplexity(model, eval_ids))
            hist.append(rec)
            if verbose:
                ev = f"  验证损失 {rec['eval_loss']:.3f}" if "eval_loss" in rec else ""
                print(f"    第 {step + 1:5d} 步  训练损失 {rec['train_loss']:.3f}{ev}  ({time.time() - t0:.0f} s)")
    model.eval()
    return hist


# ---------------------------------------------------------------------------
# 3. 标准模型：有缓存就读，没有就训练
# ---------------------------------------------------------------------------

# 每个标准模型的"配方"。改动任何一项都会得到一个新的缓存文件（文件名里含配方的哈希）。
STANDARD_MODELS = {
    "base": dict(world_seed=0, corpus_seed=0, block_size=64, n_layer=2, n_head=4, d_model=96,
                 steps=3000, batch_size=32, lr=3e-3, seed=0),
}


def _recipe_hash(recipe, train_text):
    """
    把"配方 + 训练语料的内容"转成稳定的短哈希，用作缓存文件名的一部分。

    【为什么要把语料也算进去】只看配方的话，改了 pretrain_world.py 里的模板之后，配方没变，
        就会读到用旧语料训练的旧模型。把语料文本也放进哈希，语料一变，文件名就变，自动重新训练。
    【Python 小知识】json.dumps(sort_keys=True) 保证字典键的顺序固定，同样的配方永远得到同样的字符串。
    """
    blob = json.dumps(recipe, sort_keys=True) + "\n" + train_text
    return hashlib.sha256(blob.encode()).hexdigest()[:10]


def build_data(recipe):
    """按配方生成世界、语料、分词器、训练 / 验证 token 张量，以及训练语料的原文（用来算缓存哈希）。"""
    world = build_world(seed=recipe["world_seed"])
    train_text = corpus_text(build_corpus(world, seed=recipe["corpus_seed"]))
    eval_text = corpus_text(build_corpus(world, seed=recipe["corpus_seed"] + 1, scale=0.2))
    tok = WordTokenizer(train_text)
    train_ids = torch.tensor(tok.encode(train_text), dtype=torch.long)
    eval_ids = torch.tensor(tok.encode(eval_text), dtype=torch.long)
    return world, tok, train_ids, eval_ids, train_text


def get_model(name="base", quick=False, verbose=True):
    """
    取一个标准模型。

    【参数】name：STANDARD_MODELS 里的名字；quick：True 时只训练 60 步（smoke test，不读写缓存）；
        verbose：是否打印训练过程
    【返回】(model, tok, info)：info 里有 world、配方、训练历史、是否来自缓存、训练用时
    【边界】缓存文件在 code/lecture01/checkpoints/ 下；删掉它就会重新训练（约 50 秒）。
        缓存里存了词表，读取时会检查和重新生成的词表是否一致，不一致就重新训练。
    【Python 小知识】torch.load(..., weights_only=True) 只允许读取张量和基本类型（列表、字典、字符串），
        不会执行文件里的任意代码——读取别人给的模型文件时应该总是这样做。
    """
    torch.set_num_threads(1)
    recipe = dict(STANDARD_MODELS[name])
    if quick:
        recipe["steps"] = 60
    world, tok, train_ids, eval_ids, train_text = build_data(recipe)
    cfg = GPTConfig(vocab_size=len(tok), block_size=recipe["block_size"], n_layer=recipe["n_layer"],
                    n_head=recipe["n_head"], d_model=recipe["d_model"])
    path = os.path.join(CKPT_DIR, f"{name}-{_recipe_hash(recipe, train_text)}.pt")
    info = {"world": world, "recipe": recipe, "path": path, "cached": False}
    if not quick and os.path.exists(path):
        ckpt = torch.load(path, weights_only=True)
        if ckpt["vocab"] == tok.itos:
            torch.manual_seed(recipe["seed"])
            model = TinyGPT(cfg)
            model.load_state_dict(ckpt["state_dict"])
            model.eval()
            info.update(cached=True, history=ckpt["history"], train_seconds=ckpt["train_seconds"])
            if verbose:
                print(f"  读取缓存的模型：{os.path.relpath(path, HERE)}")
            return model, tok, info
    torch.manual_seed(recipe["seed"])
    model = TinyGPT(cfg)
    if verbose:
        print(f"  训练迷你 GPT \"{name}\"：{model.n_params():,} 个参数，{len(train_ids):,} 个训练 token，"
              f"{recipe['steps']} 步 × {recipe['batch_size']} × {recipe['block_size']}")
    t0 = time.time()
    hist = train_model(model, train_ids, steps=recipe["steps"], batch_size=recipe["batch_size"],
                       lr=recipe["lr"], seed=recipe["seed"], eval_ids=eval_ids,
                       log_every=max(1, recipe["steps"] // 6), verbose=verbose)
    secs = time.time() - t0
    info.update(history=hist, train_seconds=secs)
    if not quick:
        os.makedirs(CKPT_DIR, exist_ok=True)
        torch.save({"state_dict": model.state_dict(), "vocab": tok.itos, "history": hist,
                    "train_seconds": secs, "recipe": recipe, "config": asdict(cfg)}, path)
        if verbose:
            print(f"  已保存到 {os.path.relpath(path, HERE)}（训练用时 {secs:.0f} 秒）")
    return model, tok, info


# ---------------------------------------------------------------------------
# 4. 使用模型：下一个 token 的分布、打分、续写、困惑度
# ---------------------------------------------------------------------------

def _encode_prompt(model, tok, prompt):
    """把提示编码成 (1, T) 张量；超过上下文窗口时只保留最后 block_size 个 token（左截断）。"""
    ids = tok.encode(prompt)[-model.cfg.block_size:]
    return torch.tensor([ids], dtype=torch.long)


@torch.no_grad()
def next_token_probs(model, tok, prompt):
    """
    给定提示，下一个 token 的概率分布。

    【返回】长度为 V 的一维张量（和为 1）
    【边界】提示为空字符串时没有任何 token 可以条件，会报错；调用方至少给 "<doc>"。
    【Python 小知识】@torch.no_grad() 写在函数定义上方（"装饰器"），等于整个函数体都在 no_grad 里运行。
    """
    x = _encode_prompt(model, tok, prompt)
    assert x.shape[1] > 0, "提示不能为空（至少给一个 '<doc>'）"
    logits, _, _ = model(x)
    return F.softmax(logits[0, -1], dim=-1)


def top_next(model, tok, prompt, k=5):
    """概率最大的 k 个下一 token：[(词, 概率), ...]。"""
    p = next_token_probs(model, tok, prompt)
    vals, idx = torch.topk(p, k)
    return [(tok.itos[i], v) for i, v in zip(idx.tolist(), vals.tolist())]


def score_options(model, tok, prompt, options):
    """
    几个候选词作为"下一个 token"的概率，以及它们之间归一化后的概率。

    【参数】options：候选词列表，例如 ["positive", "negative"]（每个必须是词表里的单个词）
    【返回】(原始概率列表, 归一化概率列表, 概率最大的候选词)
    【提示工程联系】这是用语言模型做分类的标准办法：不让它自由生成，而是比较几个"标签词"
        （verbalizer）的概率。第 2 讲的上下文学习实验全都建立在它上面。
    """
    p = next_token_probs(model, tok, prompt)
    raw = [p[tok.stoi[o]].item() for o in options]
    s = sum(raw)
    norm = [r / s for r in raw] if s > 0 else [1.0 / len(options)] * len(options)
    best = options[max(range(len(options)), key=lambda i: raw[i])]
    return raw, norm, best


@torch.no_grad()
def generate(model, tok, prompt, max_new_tokens=20, temperature=0.0, top_k=None,
             stop=(".", "<doc>"), seed=0):
    """
    从提示开始续写。

    【参数】temperature：0 表示贪心（每次取概率最大的 token）；> 0 时按 softmax(logits / T) 采样；
        top_k：只在概率最大的 k 个 token 里采样（None 表示不限制）；
        stop：生成到这些 token 之一就停（停止 token 本身包含在输出里）；seed：采样的随机种子
    【返回】新生成的文字（不含提示）
    【边界】提示太长时只看最后 block_size 个 token；temperature 很小（如 1e-6）时等价于贪心。
    """
    gen = torch.Generator().manual_seed(seed)
    ids = tok.encode(prompt)
    out = []
    for _ in range(max_new_tokens):
        x = torch.tensor([ids[-model.cfg.block_size:]], dtype=torch.long)
        logits = model(x)[0][0, -1]
        if temperature <= 0:
            nxt = int(torch.argmax(logits))
        else:
            logits = logits / temperature
            if top_k is not None:
                kth = torch.topk(logits, top_k).values[-1]
                logits = logits.masked_fill(logits < kth, float("-inf"))
            nxt = int(torch.multinomial(F.softmax(logits, dim=-1), 1, generator=gen))
        ids.append(nxt)
        out.append(nxt)
        if tok.itos[nxt] in stop:
            break
    return tok.decode(out)


@torch.no_grad()
def hidden_states(model, tok, text):
    """
    一段文字在最后一层的逐位置向量。

    【返回】形状 (T, d_model) 的张量，T = 文字的 token 数（超过窗口时只保留最后 block_size 个）
    【思路】因为是因果注意力，第 t 个向量只"看过"第 1…t 个 token：同样一个词 "cold"，
        在 "was cold" 和 "was not cold" 里得到的向量不一样——这就是"上下文相关的表示"。
    """
    h, _ = model.features(_encode_prompt(model, tok, text))
    return h[0]


@torch.no_grad()
def perplexity(model, ids):
    """
    在一条 token 流上的困惑度：切成不重叠的 block_size+1 长的片段，每段预测后 block_size 个 token。

    【返回】困惑度（浮点数）
    【边界】每段开头的几个 token 上下文很短，所以这个数字比"每个 token 都有完整窗口"时略高一点。
    """
    model.eval()
    T = model.cfg.block_size
    nll, n = 0.0, 0
    for i in range(0, len(ids) - 1, T):
        chunk = ids[i:i + T + 1]
        if len(chunk) < 2:
            break
        x, y = chunk[:-1].unsqueeze(0), chunk[1:].unsqueeze(0)
        _, loss, _ = model(x, y)
        nll += loss.item() * y.numel()
        n += y.numel()
    return math.exp(nll / n)


if __name__ == "__main__":
    quick = "--quick" in sys.argv
    model, tok, info = get_model("base", quick=quick)
    print(f"  参数 {model.n_params():,}；来自缓存：{info['cached']}")
    for prompt in ["<doc> faq : q : what is the capital of " + info["world"].countries[-1].name + " ? a :",
                   "<doc> reviews : review : the soup was cold . rating :",
                   "<doc> story : the queen"]:
        print(f"  {prompt!r} → {generate(model, tok, prompt)!r}")