"""第 20 章：可训练的微型解码器，以及位置编码的数值演示。"""

from math import sqrt
import torch
from torch import nn


class CausalAttention(nn.Module):
    def __init__(self, width, heads):
        super().__init__()
        if width % heads:
            raise ValueError("width 必须能被 heads 整除")
        self.heads = heads
        self.qkv = nn.Linear(width, 3*width, bias=False)
        self.output = nn.Linear(width, width, bias=False)

    def forward(self, x):
        batch, length, width = x.shape
        head_width = width//self.heads
        projected = self.qkv(x).chunk(3, dim=-1)
        query, key, value = [
            part.reshape(batch, length, self.heads, head_width).transpose(1, 2)
            for part in projected
        ]
        scores = query @ key.transpose(-2, -1) / sqrt(head_width)
        allowed = torch.ones(length, length, dtype=torch.bool, device=x.device).tril()
        weights = scores.masked_fill(~allowed, float("-inf")).softmax(-1)
        mixed = (weights @ value).transpose(1, 2).contiguous().reshape(batch, length, width)
        return self.output(mixed)


class Block(nn.Module):
    def __init__(self, width, heads):
        super().__init__()
        self.norm1, self.norm2 = nn.LayerNorm(width), nn.LayerNorm(width)
        self.attention = CausalAttention(width, heads)
        self.feedforward = nn.Sequential(nn.Linear(width, 4*width), nn.GELU(),
                                         nn.Linear(4*width, width))

    def forward(self, x):
        x = x+self.attention(self.norm1(x))
        return x+self.feedforward(self.norm2(x))


class TinyTransformer(nn.Module):
    def __init__(self, vocab_size, max_context=32, width=32, heads=4, layers=2):
        super().__init__()
        self.config = dict(vocab_size=vocab_size, max_context=max_context,
                           width=width, heads=heads, layers=layers)
        self.tokens = nn.Embedding(vocab_size, width)
        self.positions = nn.Embedding(max_context, width)
        self.blocks = nn.Sequential(*[Block(width, heads) for _ in range(layers)])
        self.final_norm = nn.LayerNorm(width)
        self.output = nn.Linear(width, vocab_size, bias=False)
        self.apply(self.initialize)

    @staticmethod
    def initialize(module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                nn.init.zeros_(module.bias)

    def forward(self, ids):
        length = ids.shape[1]
        if not 1 <= length <= self.config["max_context"]:
            raise ValueError("输入长度必须在 1 与 max_context 之间")
        positions = torch.arange(length, device=ids.device)
        x = self.tokens(ids)+self.positions(positions)
        return self.output(self.final_norm(self.blocks(x)))


def sinusoidal(length, width):
    if width % 2:
        raise ValueError("本演示要求偶数维度")
    positions = torch.arange(length, dtype=torch.float64)[:, None]
    frequencies = 10000.0**(-torch.arange(0, width, 2, dtype=torch.float64)/width)
    angles = positions*frequencies
    result = torch.empty(length, width, dtype=torch.float64)
    result[:, 0::2], result[:, 1::2] = angles.sin(), angles.cos()
    return result


def rotate_pairs(vector, position, base=10000.0):
    if vector.shape[-1] % 2:
        raise ValueError("旋转维度必须是偶数")
    width = vector.shape[-1]
    frequencies = base**(-torch.arange(0, width, 2, dtype=vector.dtype,
                                      device=vector.device)/width)
    angles = position*frequencies
    even, odd = vector[..., 0::2], vector[..., 1::2]
    rotated = torch.stack([even*angles.cos()-odd*angles.sin(),
                           even*angles.sin()+odd*angles.cos()], dim=-1)
    return rotated.flatten(-2)


def main():
    torch.set_num_threads(1)
    torch.manual_seed(2026)
    model = TinyTransformer(vocab_size=20, max_context=8, width=16, heads=4, layers=2).double()
    model.eval()
    ids = torch.tensor([[1, 3, 5, 7], [1, 4, 6, 8]])
    with torch.no_grad():
        original = model(ids)
        changed_ids = ids.clone()
        changed_ids[:, 3] = 9
        changed = model(changed_ids)
    print(f"输入 ID 形状：{tuple(ids.shape)}；输出分数形状：{tuple(original.shape)}")
    print(f"可学习参数数：{sum(p.numel() for p in model.parameters())}")
    print(f"改变末项后，前三个位置的输出不变："
          f"{torch.equal(original[:, :3], changed[:, :3])}")
    print(f"正弦位置编码第 0 行：{sinusoidal(3, 4)[0].tolist()}")
    query = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float64)
    key = torch.tensor([2.0, 0.0, 1.0, -1.0], dtype=torch.float64)
    left = torch.dot(rotate_pairs(query, 2), rotate_pairs(key, 5))
    right = torch.dot(query, rotate_pairs(key, 3))
    print(f"旋转前后长度相同：{torch.allclose(query.norm(), rotate_pairs(query, 2).norm())}")
    print(f"位置 2 与 5 的点积={left:.6f}；相对位置 3 的点积={right:.6f}")
    print(f"相对位置恒等式成立：{torch.allclose(left, right, atol=1e-12)}")


if __name__ == "__main__":
    main()