"""
modern_transformer.py —— 第 4 讲：一个"每个零件都能换"的 decoder-only Transformer
=====================================================================

【这个文件要回答的问题】
    第 2 讲实现的是"教科书版"的 Transformer。从 2017 年的原版到今天的 Llama 3，
    slides 把改动归到四个地方，这里每个都做成一个开关：

        前馈层 ffn      ："relu"（原版）、"gelu"（GPT-2 / BERT）、"swiglu"（PaLM / Llama）、"geglu"（T5 v1.1）
        注意力 n_kv_heads：= n_heads 是多头（MHA），= 1 是多查询（MQA），介于两者之间是分组查询（GQA）
        归一化 norm / norm_pos："layernorm" 或 "rmsnorm"；"post"（原版）或 "pre"（GPT-2 起）
        位置编码 pos    ："sinusoidal"（原版）、"learned"（GPT-2 / BERT）、"t5bias"（T5 的相对位置偏置）、
                         "rope"（RoFormer / Llama）、"none"（NoPE）
        偏置 bias       ：线性层要不要偏置（Llama 不要）

    另外实现了 KV cache（generate 时逐个 token 解码，复用前面位置的 K、V），用来核对 GQA 省下的显存。

【怎样使用】
    from modern_transformer import ModelConfig, ModernLM, LLAMA_STYLE, VANILLA_2017
    cfg = ModelConfig(vocab=100, **LLAMA_STYLE)
    model = ModernLM(cfg)
    logits = model(ids)                                 # ids: (B, T) → (B, T, V)
    logits, cache = model(ids, cache=[])                # 带 KV cache 的前向（见 decode_step）

【约定】和第 2 讲相同：嵌入初始化为 N(0, 1/d)、乘以 √d；输入输出共享嵌入；因果掩码；
        于是 ModelConfig 取第 2 讲的设置时（gelu、layernorm、pre、sinusoidal、bias、n_kv_heads = n_heads），
        可以直接载入第 2 讲训练好的参数（load_from_lecture02），输出完全相同。

【LLM 联系】
    slides 的总结：这些改动大多是为了**效率**和**大规模训练的稳定性**，而且"简单往往获胜"。
    Llama 3 的选择：SwiGLU、GQA、pre-norm + RMSNorm、RoPE。
"""

import math
from dataclasses import dataclass

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


@dataclass
class ModelConfig:
    vocab: int = 100
    d_model: int = 64
    n_heads: int = 4
    n_kv_heads: int = 4
    n_layers: int = 2
    d_ff: int = 256
    ffn: str = "gelu"
    norm: str = "layernorm"
    norm_pos: str = "pre"
    pos: str = "sinusoidal"
    bias: bool = True
    max_len: int = 512
    rope_base: float = 10000.0
    t5_buckets: int = 32
    t5_max_distance: int = 128


# 两套"配方"（d = 64、4 个头、2 层时，两者的总参数量都约是 10.6 万）。
# 门控 FFN 有三个矩阵，中间层通常取普通 FFN 的 2/3（4d × 2/3 ≈ 171）；这里取 196，
# 把 GQA（K、V 只有一半的头）和去掉偏置省下的参数补回 FFN，让两套配方的总参数量几乎相同。
VANILLA_2017 = dict(ffn="relu", norm="layernorm", norm_pos="post", pos="sinusoidal", bias=True,
                    n_heads=4, n_kv_heads=4, d_ff=256)
LLAMA_STYLE = dict(ffn="swiglu", norm="rmsnorm", norm_pos="pre", pos="rope", bias=False,
                   n_heads=4, n_kv_heads=2, d_ff=196)


# =====================================================================
# 1. 归一化
# =====================================================================

class RMSNorm(nn.Module):
    """
    RMSNorm（Zhang & Sennrich 2019）：y = x / sqrt(mean(x²) + eps) × g。

    【和 LayerNorm 的区别】LayerNorm 先减均值（"重新居中"）、再除以标准差、再乘缩放 g、加偏置 b；
            RMSNorm 只做"按均方根缩放"再乘 g：不减均值、没有偏置。
            少一次求均值、少一组参数，计算更省；实验上效果和 LayerNorm 相当甚至更好。
    【边界】eps 防止全零向量除以 0。数值上与 PyTorch 2.4+ 的 nn.RMSNorm 一致（02 号脚本核对）。
    """

    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight


def make_norm(cfg):
    return nn.LayerNorm(cfg.d_model) if cfg.norm == "layernorm" else RMSNorm(cfg.d_model)


# =====================================================================
# 2. 前馈层
# =====================================================================

class FeedForward(nn.Module):
    """
    【普通 FFN】  down(φ(up(x)))，φ = ReLU / GELU；两个矩阵，参数 2·d·d_ff。
    【门控 FFN】  down(φ(gate(x)) ⊙ up(x))，φ = SiLU（SwiGLU）或 GELU（GEGLU）；三个矩阵，参数 3·d·d_ff。
            ⊙ 是逐元素相乘：一条支路决定"开多大的门"，另一条支路提供"内容"——
            输出对输入是二次的（"双线性"），表达能力更强。为了和普通 FFN 参数量相同，
            门控 FFN 的 d_ff 取普通 FFN 的 2/3（Shazeer 2020）。
    """

    def __init__(self, cfg):
        super().__init__()
        self.kind = cfg.ffn
        self.up = nn.Linear(cfg.d_model, cfg.d_ff, bias=cfg.bias)
        self.down = nn.Linear(cfg.d_ff, cfg.d_model, bias=cfg.bias)
        if self.kind in ("swiglu", "geglu"):
            self.gate = nn.Linear(cfg.d_model, cfg.d_ff, bias=cfg.bias)

    def forward(self, x):
        if self.kind == "relu":
            return self.down(F.relu(self.up(x)))
        if self.kind == "gelu":
            return self.down(F.gelu(self.up(x)))
        act = F.silu if self.kind == "swiglu" else F.gelu
        return self.down(act(self.gate(x)) * self.up(x))


# =====================================================================
# 3. 位置编码
# =====================================================================

def sinusoidal_positions(T, d, offset=0):
    """第 2 讲的正弦位置编码（从 offset 开始的 T 个位置）。【返回】(T, d)。"""
    pos = torch.arange(offset, offset + T, dtype=torch.float32)[:, None]
    freq = torch.exp(-math.log(10000.0) * torch.arange(0, d, 2) / d)
    pe = torch.zeros(T, d)
    pe[:, 0::2] = torch.sin(pos * freq)
    pe[:, 1::2] = torch.cos(pos * freq)
    return pe


def rope_angles(positions, head_dim, base=10000.0):
    """
    RoPE 的旋转角：位置 m、第 i 对维度的角度 = m · θ_i，θ_i = base^(-2i / head_dim)。

    【参数】positions：(T,) 的整数张量；head_dim：每个头的维度（必须是偶数）。
    【返回】(cos, sin)，形状都是 (T, head_dim / 2)。
    【直觉】每对维度是一根以 θ_i 的角速度旋转的"指针"：低维转得快，高维转得慢——和正弦位置编码的频率一模一样，
            区别在于 RoPE 用它去**旋转** q、k，而不是**加**到嵌入上。
    """
    theta = base ** (-torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
    ang = positions.float()[:, None] * theta[None, :]
    return ang.cos(), ang.sin()


def apply_rope(x, cos, sin):
    """
    把 x 的每一对相邻维度 (2i, 2i+1) 旋转对应的角度：
        [x_{2i}, x_{2i+1}] → [x_{2i} cos − x_{2i+1} sin,  x_{2i} sin + x_{2i+1} cos]

    【参数】x：(..., T, head_dim)；cos、sin：(T, head_dim/2)。
    【返回】与 x 同形状。
    【为什么注意力分数只依赖相对位置】把 q 转 mθ、k 转 nθ，二维旋转矩阵满足 R(mθ)ᵀ R(nθ) = R((n−m)θ)，
            所以 (R_m q)·(R_n k) = qᵀ R_{n−m} k 只依赖 n − m。04 号脚本数值验证。
    【边界】这里用"相邻两维一对"的写法（RoFormer 论文、slides 的写法）；Llama 的代码用"前一半和后一半配对"，
            两者只是维度的排列不同，数学上等价。
    """
    x1, x2 = x[..., 0::2], x[..., 1::2]
    out = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
    return out.flatten(-2)


def t5_bucket(rel, num_buckets=32, max_distance=128):
    """
    T5 的相对位置分桶（单向 / 因果版本）：距离 rel = 查询位置 − 键位置（≥ 0）。

    【思路】近的距离每个一个桶（区分 1、2、3 个词之外很重要），远的距离按对数分桶（100 和 110 的区别不大），
            超过 max_distance 的全部落进最后一个桶——所以它对"没见过的长度"天然有一定的外推能力。
    【返回】与 rel 同形状的桶编号，范围 [0, num_buckets)。
    """
    rel = rel.clamp(min=0)
    exact = num_buckets // 2
    is_small = rel < exact
    large = exact + (torch.log(rel.float().clamp(min=1) / exact) / math.log(max_distance / exact)
                     * (num_buckets - exact)).long()
    large = large.clamp(max=num_buckets - 1)
    return torch.where(is_small, rel, large)


# =====================================================================
# 4. 注意力（MHA / GQA / MQA，可选 RoPE、T5 偏置、KV cache）
# =====================================================================

class Attention(nn.Module):
    """
    因果自注意力，支持分组查询。

    【参数】n_heads 个查询头，n_kv_heads 个键值头；n_heads 必须是 n_kv_heads 的整数倍。
            每 n_heads / n_kv_heads 个查询头共用一组 K、V。
    【前向的参数】
        x：(B, T, d)；offset：这 T 个 token 在序列里的起始位置（有 KV cache 时 > 0）；
        cache：None，或者一个字典 {"k": (B, kv, T_past, hd), "v": ...}，会被原地更新；
        bias：T5 相对位置偏置，(h, T, T_total)，或 None。
    【返回】(B, T, d)。
    【KV cache】缓存的是"还没展开成 n_heads 份"的 K、V——GQA 省显存就省在这里。
    """

    def __init__(self, cfg):
        super().__init__()
        assert cfg.n_heads % cfg.n_kv_heads == 0, "n_heads 必须是 n_kv_heads 的整数倍"
        self.h, self.kv, self.hd = cfg.n_heads, cfg.n_kv_heads, cfg.d_model // cfg.n_heads
        self.pos, self.base = cfg.pos, cfg.rope_base
        self.q_proj = nn.Linear(cfg.d_model, self.h * self.hd, bias=cfg.bias)
        self.k_proj = nn.Linear(cfg.d_model, self.kv * self.hd, bias=cfg.bias)
        self.v_proj = nn.Linear(cfg.d_model, self.kv * self.hd, bias=cfg.bias)
        self.o_proj = nn.Linear(self.h * self.hd, cfg.d_model, bias=cfg.bias)

    def forward(self, x, offset=0, cache=None, bias=None):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.h, self.hd).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.kv, self.hd).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.kv, self.hd).transpose(1, 2)
        if self.pos == "rope":
            cos, sin = rope_angles(torch.arange(offset, offset + T), self.hd, self.base)
            q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        if cache is not None:
            if "k" in cache:
                k = torch.cat([cache["k"], k], dim=2)
                v = torch.cat([cache["v"], v], dim=2)
            cache["k"], cache["v"] = k, v
        Tk = k.shape[2]
        rep = self.h // self.kv
        if rep > 1:                                   # 每组 K、V 复制给组内的查询头（只在计算时展开）
            k = k.repeat_interleave(rep, dim=1)
            v = v.repeat_interleave(rep, dim=1)
        scores = q @ k.transpose(-1, -2) / math.sqrt(self.hd)          # (B, h, T, Tk)
        if bias is not None:
            scores = scores + bias
        qpos = torch.arange(offset, offset + T)[:, None]
        kpos = torch.arange(Tk)[None, :]
        scores = scores.masked_fill(kpos > qpos, float("-inf"))
        out = scores.softmax(-1) @ v
        return self.o_proj(out.transpose(1, 2).reshape(B, T, self.h * self.hd))


class Block(nn.Module):
    """一个块：注意力 + FFN，各带残差；归一化放在子层之前（pre）或残差相加之后（post）。"""

    def __init__(self, cfg):
        super().__init__()
        self.post = cfg.norm_pos == "post"
        self.n1, self.attn = make_norm(cfg), Attention(cfg)
        self.n2, self.ffn = make_norm(cfg), FeedForward(cfg)

    def forward(self, x, offset=0, cache=None, bias=None):
        if self.post:
            x = self.n1(x + self.attn(x, offset, cache, bias))
            return self.n2(x + self.ffn(x))
        x = x + self.attn(self.n1(x), offset, cache, bias)
        return x + self.ffn(self.n2(x))


class ModernLM(nn.Module):
    """
    decoder-only 语言模型，零件由 ModelConfig 决定。

    【前向的参数】ids：(B, T)；pad：为了和第 2 讲的训练工具兼容而保留（右侧补齐时因果模型用不到）；
            cache：None 表示普通前向；传入一个列表（第一次传 []）则启用 KV cache，返回 (logits, cache)。
    【边界】pos="learned" 时总长度不能超过 max_len；其余位置编码对长度没有硬限制。
    """

    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        d = cfg.d_model
        self.emb = nn.Embedding(cfg.vocab, d)
        nn.init.normal_(self.emb.weight, std=d ** -0.5)
        if cfg.pos == "learned":
            self.pos_emb = nn.Embedding(cfg.max_len, d)
            nn.init.normal_(self.pos_emb.weight, std=0.02 * d ** 0.5)
        if cfg.pos == "t5bias":
            self.rel_bias = nn.Embedding(cfg.t5_buckets, cfg.n_heads)      # 所有层共用一张表（T5 的做法）
            nn.init.zeros_(self.rel_bias.weight)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
        self.norm_f = make_norm(cfg) if cfg.norm_pos == "pre" else nn.Identity()

    def _t5_bias(self, offset, T, Tk):
        rel = torch.arange(offset, offset + T)[:, None] - torch.arange(Tk)[None, :]
        b = self.rel_bias(t5_bucket(rel, self.cfg.t5_buckets, self.cfg.t5_max_distance))   # (T, Tk, h)
        return b.permute(2, 0, 1)

    def forward(self, ids, pad=None, cache=None):
        B, T = ids.shape
        offset = 0
        if cache is not None and len(cache) == len(self.blocks):
            offset = cache[0]["k"].shape[2]
        elif cache is not None:
            cache.extend({} for _ in self.blocks)
        x = self.emb(ids) * math.sqrt(self.cfg.d_model)
        if self.cfg.pos == "sinusoidal":
            x = x + sinusoidal_positions(T, self.cfg.d_model, offset)
        elif self.cfg.pos == "learned":
            x = x + self.pos_emb(torch.arange(offset, offset + T))
        bias = self._t5_bias(offset, T, offset + T) if self.cfg.pos == "t5bias" else None
        for i, blk in enumerate(self.blocks):
            x = blk(x, offset, None if cache is None else cache[i], bias)
        logits = self.norm_f(x) @ self.emb.weight.T
        return (logits, cache) if cache is not None else logits


def count_params(model):
    return sum(p.numel() for p in model.parameters())


def kv_cache_elements(cache):
    """KV cache 里一共存了多少个数（所有层的 K 和 V）。"""
    return sum(c["k"].numel() + c["v"].numel() for c in cache)


def load_from_lecture02(model, sd, kv_init="mean"):
    """
    把第 2 讲 TransformerLM（gelu、layernorm、pre、sinusoidal、带偏置）的参数载入 ModernLM。

    【参数】sd：第 2 讲模型的 state_dict；kv_init：当 ModernLM 的 n_kv_heads 少于 n_heads 时，
            怎样从多头的 K、V 投影得到分组的投影：
              "mean"   组内各头的投影矩阵取平均（Ainslie et al. 2023 的做法）
              "first"  只取组内第一个头
              "random" 重新随机初始化（相当于不继承 K、V）
    【返回】model（原地修改）。
    【思路】名字一一对应：blocks.i.self_attn.{q,k,v,o}_proj → blocks.i.attn.*，n1 → n1，n3 → n2，
            ffn.up / ffn.down → ffn.up / ffn.down，norm_f → norm_f。K、V 投影的权重形状是 (h·hd, d)，
            先 view 成 (h, hd, d)，按组合并。
    """
    cfg = model.cfg
    h, kv, hd = cfg.n_heads, cfg.n_kv_heads, cfg.d_model // cfg.n_heads
    new = {"emb.weight": sd["emb.weight"], "norm_f.weight": sd["norm_f.weight"], "norm_f.bias": sd["norm_f.bias"]}
    for i in range(cfg.n_layers):
        p = f"blocks.{i}."
        for a in ("q", "o"):
            for t in ("weight", "bias"):
                new[p + f"attn.{a}_proj.{t}"] = sd[p + f"self_attn.{a}_proj.{t}"]
        for a in ("k", "v"):
            w = sd[p + f"self_attn.{a}_proj.weight"].view(h, hd, -1)
            b = sd[p + f"self_attn.{a}_proj.bias"].view(h, hd)
            g = h // kv
            if kv_init == "mean":
                w2, b2 = w.view(kv, g, hd, -1).mean(1), b.view(kv, g, hd).mean(1)
            elif kv_init == "first":
                w2, b2 = w.view(kv, g, hd, -1)[:, 0], b.view(kv, g, hd)[:, 0]
            else:
                w2 = model.state_dict()[p + f"attn.{a}_proj.weight"].view(kv, hd, -1)
                b2 = model.state_dict()[p + f"attn.{a}_proj.bias"].view(kv, hd)
            new[p + f"attn.{a}_proj.weight"] = w2.reshape(kv * hd, -1)
            new[p + f"attn.{a}_proj.bias"] = b2.reshape(kv * hd)
        for t in ("weight", "bias"):
            new[p + f"n1.{t}"] = sd[p + f"n1.{t}"]
            new[p + f"n2.{t}"] = sd[p + f"n3.{t}"]
            new[p + f"ffn.up.{t}"] = sd[p + f"ffn.up.{t}"]
            new[p + f"ffn.down.{t}"] = sd[p + f"ffn.down.{t}"]
    model.load_state_dict(new)
    return model