"""对照组：标准的 GPT 式 Transformer（因果自注意力），结构上和 ZeroX 一一对应。

位置编码两种可选（《Attention Is All You Need》§3.5 两种都试过，结果几乎一样）：
    pos="learned"  可学习的位置向量，只有 block_size 个，处理不了更长的序列
    pos="sin"      论文用的正弦编码，任意长度都能算（论文选它就是因为"可能可以外推"）
"""
import numpy as np

from zerocg import Tensor
from zerocg import functional as F
from zerocg import nn


def sinusoidal_encoding(T, d):
    """PE(pos, 2i) = sin(pos / 10000^(2i/d))，PE(pos, 2i+1) = cos(...)。论文 §3.5。"""
    pos = np.arange(T)[:, None]
    div = 10000 ** (np.arange(0, d, 2) / d)
    pe = np.zeros((T, d), dtype=np.float32)
    pe[:, 0::2] = np.sin(pos / div)
    pe[:, 1::2] = np.cos(pos / div)
    return pe


class Block(nn.Module):
    def __init__(self, n_embd, n_head, dropout=0.0):
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = nn.CausalSelfAttention(n_embd, n_head, dropout)
        self.ln2 = nn.LayerNorm(n_embd)
        self.fc1 = nn.Linear(n_embd, 4 * n_embd)
        self.fc2 = nn.Linear(4 * n_embd, n_embd)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        return x + self.drop(self.fc2(F.gelu(self.fc1(self.ln2(x)))))


class TransformerLM(nn.Module):
    def __init__(self, vocab_size, block_size, n_embd=64, n_head=4, n_layer=2, dropout=0.0, pos="learned"):
        assert pos in ("learned", "sin")
        self.block_size = block_size
        self.pos_kind = pos
        self.tok = nn.Embedding(vocab_size, n_embd)
        if pos == "learned":
            self.pos = nn.Embedding(block_size, n_embd)
        self.blocks = [Block(n_embd, n_head, dropout) for _ in range(n_layer)]
        self.ln_f = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx):
        T = idx.shape[1]
        if self.pos_kind == "learned":
            assert T <= self.block_size, f"位置编码只有 {self.block_size} 个，处理不了 {T} 个字"
            x = self.tok(idx) + self.pos(np.arange(T))
        else:
            x = self.tok(idx) + Tensor(sinusoidal_encoding(T, self.tok.weight.shape[1]))
        for b in self.blocks:
            x = b(x)
        return self.head(self.ln_f(x))

    def state_floats(self, T):
        """推理到第 T 个字时 KV 缓存的大小（浮点数个数），随 T 线性增长。"""
        return sum(2 * T * b.ln1.weight.shape[0] for b in self.blocks)