"""
tiny_gpt.py —— 一个可配置的小 Transformer 语言模型（只用来"数参数、测显存、数 FLOPs"）
=====================================================================

【这个文件做什么】
    前置讲 P3 要回答"一个模型有多少参数、训练要多少显存、要算多少次乘加"。
    为了用真实的 PyTorch 模型核对这些公式，这里写了一个结构和 GPT-2 / Llama 一致、
    但可以做得很小的 decoder-only Transformer。
    **结构细节（注意力、残差、归一化……）在第 2、4 讲才讲**；这里你只需要知道它由
    嵌入层 + L 个"块"（每块 = 注意力 + 前馈网络 MLP）+ 输出层组成。

【可以切换的设计】（对应 mem_calc.Config 的字段）
    - mlp="gelu"（GPT-2：两个矩阵）或 "swiglu"（Llama：三个矩阵）
    - norm="layernorm"（有缩放和偏置两组参数）或 "rmsnorm"（只有缩放）
    - bias：线性层要不要偏置（GPT-2 要，Llama 不要）
    - tied：输出层是否和输入嵌入共享权重（GPT-2 共享，Llama 不共享）
    - learned_pos：是否有可学习的位置嵌入（GPT-2 有；Llama 用 RoPE，没有这部分参数）
    - n_kv_heads < n_heads：分组查询注意力（GQA，Llama 3 用；第 4 讲细讲）
    - checkpoint=True：激活检查点（反向时重算每个块的中间结果，用时间换显存）

【LLM 联系】
    注意力里显式地算出了 (B, 头数, T, T) 的注意力矩阵——这正是"上下文越长、激活显存
    按 T² 增长"的来源。FlashAttention（第 19 讲）通过分块计算避免存下这个矩阵。
"""

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint


class Attention(nn.Module):
    """多头因果自注意力（支持 GQA）。【输入/输出】形状 (B, T, d)。"""

    def __init__(self, d, n_heads, n_kv_heads, bias):
        super().__init__()
        self.h, self.kv, self.hd = n_heads, n_kv_heads, d // n_heads
        self.q = nn.Linear(d, n_heads * self.hd, bias=bias)
        self.k = nn.Linear(d, n_kv_heads * self.hd, bias=bias)
        self.v = nn.Linear(d, n_kv_heads * self.hd, bias=bias)
        self.o = nn.Linear(n_heads * self.hd, d, bias=bias)

    def forward(self, x):
        B, T, _ = x.shape
        q = self.q(x).view(B, T, self.h, self.hd).transpose(1, 2)      # (B, h, T, hd)
        k = self.k(x).view(B, T, self.kv, self.hd).transpose(1, 2)     # (B, kv, T, hd)
        v = self.v(x).view(B, T, self.kv, self.hd).transpose(1, 2)
        if self.kv != self.h:                                          # GQA：每组 K/V 被多个查询头共用
            k = k.repeat_interleave(self.h // self.kv, dim=1)
            v = v.repeat_interleave(self.h // self.kv, dim=1)
        att = (q @ k.transpose(-1, -2)) / math.sqrt(self.hd)           # (B, h, T, T)：T² 的来源
        causal = torch.ones(T, T, dtype=torch.bool).tril()
        att = att.masked_fill(~causal, float("-inf")).softmax(dim=-1)
        y = (att @ v).transpose(1, 2).reshape(B, T, self.h * self.hd)
        return self.o(y)


class MLP(nn.Module):
    def __init__(self, d, d_ff, kind, bias):
        super().__init__()
        self.kind = kind
        self.w1 = nn.Linear(d, d_ff, bias=bias)
        if kind == "swiglu":
            self.w3 = nn.Linear(d, d_ff, bias=bias)                    # SwiGLU 多一个"门"矩阵
        self.w2 = nn.Linear(d_ff, d, bias=bias)

    def forward(self, x):
        if self.kind == "swiglu":
            return self.w2(F.silu(self.w1(x)) * self.w3(x))
        return self.w2(F.gelu(self.w1(x)))


def make_norm(kind, d):
    return nn.LayerNorm(d) if kind == "layernorm" else nn.RMSNorm(d)


class Block(nn.Module):
    """一个 Transformer 块：x + 注意力(归一化(x))，再 + MLP(归一化(x))（pre-norm，第 4 讲解释）。"""

    def __init__(self, cfg):
        super().__init__()
        self.n1 = make_norm(cfg.norm, cfg.d_model)
        self.attn = Attention(cfg.d_model, cfg.n_heads, cfg.n_kv_heads, cfg.bias)
        self.n2 = make_norm(cfg.norm, cfg.d_model)
        self.mlp = MLP(cfg.d_model, cfg.d_ff, cfg.mlp, cfg.bias)

    def forward(self, x):
        x = x + self.attn(self.n1(x))
        return x + self.mlp(self.n2(x))


class TinyGPT(nn.Module):
    """
    【参数】cfg：mem_calc.Config；checkpoint：是否对每个块使用激活检查点。
    【输入】token id，形状 (B, T)，T <= cfg.max_ctx。【输出】logits，形状 (B, T, V)。
    """

    def __init__(self, cfg, checkpoint=False):
        super().__init__()
        self.cfg, self.use_ckpt = cfg, checkpoint
        self.emb = nn.Embedding(cfg.vocab, cfg.d_model)
        self.pos = nn.Embedding(cfg.max_ctx, cfg.d_model) if cfg.learned_pos else None
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
        self.norm_f = make_norm(cfg.norm, cfg.d_model)
        self.head = None if cfg.tied else nn.Linear(cfg.d_model, cfg.vocab, bias=False)

    def forward(self, ids):
        x = self.emb(ids)
        if self.pos is not None:
            x = x + self.pos(torch.arange(ids.shape[1]))
        for blk in self.blocks:
            # 激活检查点：前向时不保存块内部的中间结果，反向时再算一遍
            x = checkpoint(blk, x, use_reentrant=False) if self.use_ckpt else blk(x)
        x = self.norm_f(x)
        return x @ self.emb.weight.T if self.head is None else self.head(x)