"""第 8—15 课。保持各个 GPT 部件可见，不依赖现成的语言模型。"""
import math
from dataclasses import asdict, dataclass
import torch
from torch import nn
from torch.nn import functional as F


@dataclass
class ModelConfig:
    vocab_size: int
    context: int = 64
    width: int = 256
    heads: int = 4
    layers: int = 4
    dropout: float = 0.1

    def __post_init__(self):
        if min(self.vocab_size, self.context, self.width, self.heads, self.layers) < 1:
            raise ValueError("模型尺寸必须是正整数。")
        if self.width % self.heads: raise ValueError("表示宽度必须能被注意力头数整除。")
        if not 0 <= self.dropout < 1: raise ValueError("随机丢弃比例必须大于等于 0 且小于 1。")


class CausalAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.heads = config.heads
        self.qkv = nn.Linear(config.width, 3 * config.width)
        self.project = nn.Linear(config.width, config.width)
        self.dropout = config.dropout
        self.residual_dropout = nn.Dropout(config.dropout)

    def forward(self, x, inspect=False):
        batch, length, width = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        def heads(t):
            return t.view(batch, length, self.heads, width // self.heads).transpose(1, 2)
        q, k, v = map(heads, (q, k, v))
        if inspect:
            # 教学路径：逐项展开。mask 的上三角是尚未出现的字。
            score = q @ k.transpose(-2, -1) / math.sqrt(k.shape[-1])
            forbidden = torch.ones(length, length, device=x.device, dtype=torch.bool).triu(1)
            weights = score.masked_fill(forbidden, float("-inf")).softmax(dim=-1)
            values = weights @ v
        else:
            # 同样的运算交给 PyTorch 的优化实现；训练时按配置使用 dropout。
            values = F.scaled_dot_product_attention(q, k, v, is_causal=True,
                        dropout_p=self.dropout if self.training else 0.0)
        merged = values.transpose(1, 2).contiguous().view(batch, length, width)
        out = self.residual_dropout(self.project(merged))
        return (out, weights) if inspect else out


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.norm_attention = nn.LayerNorm(config.width)
        self.attention = CausalAttention(config)
        self.norm_feedforward = nn.LayerNorm(config.width)
        self.feedforward = nn.Sequential(nn.Linear(config.width, 4*config.width), nn.GELU(),
                              nn.Linear(4*config.width, config.width), nn.Dropout(config.dropout))

    def forward(self, x):
        x = x + self.attention(self.norm_attention(x))
        return x + self.feedforward(self.norm_feedforward(x))


class PoetryGPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.characters = nn.Embedding(config.vocab_size, config.width)
        self.positions = nn.Embedding(config.context, config.width)
        self.dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.layers)])
        self.final_norm = nn.LayerNorm(config.width)
        self.output = nn.Linear(config.width, config.vocab_size, bias=False)
        self.apply(self._initialize)
        # 输入与输出共用同一份字表参数，减小模型。
        self.output.weight = self.characters.weight

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

    def forward(self, tokens, targets=None):
        _, length = tokens.shape
        if length > self.config.context: raise ValueError("输入超过模型的前文长度。")
        positions = torch.arange(length, device=tokens.device)
        hidden = self.dropout(self.characters(tokens) + self.positions(positions))
        for block in self.blocks: hidden = block(hidden)
        logits = self.output(self.final_norm(hidden))
        loss = None if targets is None else F.cross_entropy(logits.reshape(-1, logits.size(-1)),
                                                 targets.reshape(-1), ignore_index=-100)
        return logits, loss

    def parameter_count(self):
        return sum(p.numel() for p in self.parameters())