"""
YonetixAI - 从零训练 Decoder-only Transformer

架构：GPT-like + RoPE + SwiGLU + Pre-LN
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
from typing import Optional


@dataclass
class ModelConfig:
    vocab_size: int = 151936
    hidden_dim: int = 768
    num_layers: int = 12
    num_heads: int = 12
    ffn_hidden_dim: int = 3072
    max_seq_len: int = 512
    dropout: float = 0.1
    layer_norm_eps: float = 1e-5
    tie_word_embeddings: bool = True
    init_std: float = 0.02


class RotaryEmbedding(nn.Module):
    """旋转位置编码 RoPE"""
    def __init__(self, dim: int, max_seq_len: int = 512, theta: float = 10000.0):
        super().__init__()
        self.dim = dim
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, x: torch.Tensor, pos: Optional[torch.Tensor] = None):
        seq_len = x.size(1)
        if pos is None:
            pos = torch.arange(seq_len, device=x.device, dtype=torch.float32)
        inv_freq = self.inv_freq[None, :]  # [1, dim/2]
        pos = pos[:, None]  # [seq, 1]
        freqs = pos * inv_freq  # [seq, dim/2]
        # 扩展到全维度
        freqs = torch.cat([freqs, freqs], dim=-1)  # [seq, dim]
        return freqs

def apply_rotary(x: torch.Tensor, freqs: torch.Tensor):
    """应用 RoPE 到 Q 或 K
    x: [batch, heads, seq, dim]
    freqs: [seq, dim]
    """
    # 分离实部和虚部
    x1 = x[..., ::2]    # 偶数位
    x2 = x[..., 1::2]   # 奇数位
    cos = freqs.cos()[None, None, :, :]  # [1, 1, seq, dim]
    sin = freqs.sin()[None, None, :, :]
    # 确保维度匹配
    cos = cos[..., ::2]  # [1, 1, seq, dim/2]
    sin = sin[..., ::2]
    x_rot = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
    return x_rot


class Attention(nn.Module):
    """多头注意力 + RoPE"""
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.num_heads = config.num_heads
        self.head_dim = config.hidden_dim // config.num_heads
        self.hidden_dim = config.hidden_dim

        self.q_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)
        self.o_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)
        self.dropout = nn.Dropout(config.dropout)

        self.rotary = RotaryEmbedding(self.head_dim, config.max_seq_len)

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None):
        B, S, D = x.shape

        q = self.q_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)

        # RoPE
        freqs = self.rotary(x)  # [S, head_dim, 1]
        q = apply_rotary(q, freqs)
        k = apply_rotary(k, freqs)

        # Flash Attention
        attn = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.dropout.p if self.training else 0.0)

        attn = attn.transpose(1, 2).contiguous().view(B, S, D)
        return self.o_proj(attn)


class SwiGLU(nn.Module):
    """SwiGLU 激活函数"""
    def __init__(self, hidden_dim: int, ffn_hidden_dim: int):
        super().__init__()
        self.gate = nn.Linear(hidden_dim, ffn_hidden_dim, bias=False)
        self.up = nn.Linear(hidden_dim, ffn_hidden_dim, bias=False)
        self.down = nn.Linear(ffn_hidden_dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor):
        return self.down(F.silu(self.gate(x)) * self.up(x))


class TransformerBlock(nn.Module):
    """Transformer 层：Pre-LN + Attention + SwiGLU"""
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.attn_norm = nn.LayerNorm(config.hidden_dim, eps=config.layer_norm_eps)
        self.attn = Attention(config)
        self.ffn_norm = nn.LayerNorm(config.hidden_dim, eps=config.layer_norm_eps)
        self.ffn = SwiGLU(config.hidden_dim, config.ffn_hidden_dim)

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None):
        x = x + self.attn(self.attn_norm(x), mask)
        x = x + self.ffn(self.ffn_norm(x))
        return x


class YonetixTransformer(nn.Module):
    """完整 Decoder-only Transformer"""
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config

        self.embed = nn.Embedding(config.vocab_size, config.hidden_dim)
        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_layers)])
        self.final_norm = nn.LayerNorm(config.hidden_dim, eps=config.layer_norm_eps)
        self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)

        # 绑定 embedding 和 lm_head 权重
        if config.tie_word_embeddings:
            self.lm_head.weight = self.embed.weight

        # 初始化
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.init_std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.init_std)

    def _create_causal_mask(self, seq_len: int, device: torch.device) -> torch.Tensor:
        """创建因果注意力掩码"""
        mask = torch.triu(torch.full((seq_len, seq_len), float('-inf'), device=device), diagonal=1)
        return mask[None, None, :, :]  # [1, 1, S, S]

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        """前向传播，返回 logits"""
        B, S = input_ids.shape
        h = self.embed(input_ids) * math.sqrt(self.config.hidden_dim)
        mask = self._create_causal_mask(S, input_ids.device)

        for layer in self.layers:
            h = layer(h, mask)

        h = self.final_norm(h)
        logits = self.lm_head(h)
        return logits

    @torch.no_grad()
    def generate(self, input_ids: torch.Tensor, max_new_tokens: int = 128,
                 temperature: float = 0.7, top_k: int = 50, top_p: float = 0.9):
        """自回归生成"""
        self.eval()
        for _ in range(max_new_tokens):
            # 截断到 max_seq_len
            if input_ids.size(1) > self.config.max_seq_len:
                input_ids = input_ids[:, -self.config.max_seq_len:]

            logits = self.forward(input_ids)
            logits = logits[:, -1, :] / temperature

            # Top-K
            if top_k > 0:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, -1:]] = float('-inf')

            # Top-P (nucleus)
            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                sorted_indices_to_remove = cum_probs > top_p
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                sorted_indices_to_remove[..., 0] = 0
                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                logits[indices_to_remove] = float('-inf')

            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, next_token], dim=1)

        return input_ids


def count_params(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters())


if __name__ == "__main__":
    config = ModelConfig()
    model = YonetixTransformer(config)
    print(f"参数量: {count_params(model) / 1e6:.2f}M")
    x = torch.randint(0, 100, (2, 32))
    out = model(x)
    print(f"输入: {x.shape}, 输出: {out.shape}")