#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
XiaLi-Next 统一架构 v1.0
=========================
同一套架构，从 20M 到 7B+ 通用
核心设计：RoPE + SwiGLU + Pre-RMSNorm + GQA + Weight Tying

用法:
  from xiaoli_next import XiaLiConfig, XiaLiModel
  
  # 20M 超小模型
  cfg = XiaLiConfig.preset('20m')
  model = XiaLiModel(cfg)
  
  # 1B 中等模型  
  cfg = XiaLiConfig.preset('1b')
  model = XiaLiModel(cfg)
  
  # 7B 大模型
  cfg = XiaLiConfig.preset('7b')
  model = XiaLiModel(cfg)
"""

import math
import json
from typing import Optional, Tuple

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


# ════════════════════════════════════════════════════════════
# 1. 配置系统
# ════════════════════════════════════════════════════════════

class XiaLiConfig:
    """统一配置 - 通过 preset 快速切换规模"""
    
    def __init__(self):
        # 架构
        self.d_model: int = 320
        self.n_layers: int = 8
        self.n_heads: int = 8
        self.n_kv_heads: int = 4       # GQA: KV头数（=n_heads 就是MHA，=1 就是MQA）
        self.d_ff: int = 0             # 0=自动 = int(d_model * 8/3 * 2)  SwiGLU需要额外gate
        self.max_seq_len: int = 2048   # 最大序列长度
        self.dropout: float = 0.0      # 大模型dropout=0
        self.rope_theta: float = 10000.0  # RoPE base frequency
        self.tie_weights: bool = True  # Weight Tying
        self.vocab_size: int = 20007
        
        # 训练
        self.norm_eps: float = 1e-5
        
    @classmethod
    def preset(cls, name: str) -> 'XiaLiConfig':
        """按规模取配置"""
        cfgs = {
            # ── 超小模型 ──
            '20m': dict(d_model=288, n_layers=6, n_heads=6, n_kv_heads=3, max_seq_len=1024),
            '26m': dict(d_model=320, n_layers=8, n_heads=8, n_kv_heads=4, max_seq_len=2048),
            # ── 百M级 ──
            '100m': dict(d_model=512, n_layers=12, n_heads=8, n_kv_heads=4, max_seq_len=2048),
            '300m': dict(d_model=768, n_layers=16, n_heads=12, n_kv_heads=4, max_seq_len=4096),
            # ── 十亿级 ──
            '500m': dict(d_model=1024, n_layers=24, n_heads=16, n_kv_heads=8, max_seq_len=8192),
            '1b':   dict(d_model=1536, n_layers=24, n_heads=16, n_kv_heads=8, max_seq_len=8192),
            '3b':   dict(d_model=2560, n_layers=32, n_heads=24, n_kv_heads=8, max_seq_len=8192),
            # ── 十亿级+ ──
            '7b':   dict(d_model=4096, n_layers=32, n_heads=32, n_kv_heads=8, max_seq_len=8192, rope_theta=500000.0),
            '13b':  dict(d_model=5120, n_layers=40, n_heads=40, n_kv_heads=10, max_seq_len=8192, rope_theta=500000.0),
            '70b':  dict(d_model=8192, n_layers=80, n_heads=64, n_kv_heads=8, max_seq_len=8192, rope_theta=500000.0),
        }
        cfg = cls()
        if name in cfgs:
            for k, v in cfgs[name].items():
                setattr(cfg, k, v)
        if cfg.d_ff == 0:
            # SwiGLU 标准: 8/3 * d_model（参数量等价于 4x 标准 FFN）
            cfg.d_ff = int(cfg.d_model * 8 / 3)
        return cfg


# ════════════════════════════════════════════════════════════
# 2. RMSNorm
# ════════════════════════════════════════════════════════════

class XiaLiRMSNorm(nn.Module):
    """RMS Layer Normalization"""
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x / rms * self.weight


# ════════════════════════════════════════════════════════════
# 3. RoPE（旋转位置编码）
# ════════════════════════════════════════════════════════════

def precompute_rope_freqs(dim: int, max_len: int, theta: float = 10000.0) -> Tuple[torch.Tensor, torch.Tensor]:
    """预计算 RoPE 频率"""
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
    t = torch.arange(max_len, dtype=torch.float32)
    angles = torch.outer(t, freqs)  # [max_len, dim//2]
    cos = angles.cos().unsqueeze(0).unsqueeze(2)  # [1, max_len, 1, dim//2]
    sin = angles.sin().unsqueeze(0).unsqueeze(2)
    return cos, sin

def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """对 x 应用 RoPE。x: [batch, seq, n_heads, head_dim]"""
    head_dim = x.size(-1)
    # cos/sin 是每对维度一个值 (head_dim//2)，需要扩张到 head_dim
    cos = torch.repeat_interleave(cos, 2, dim=-1)  # [..., d//2] -> [..., d]
    sin = torch.repeat_interleave(sin, 2, dim=-1)
    x1 = x[..., :head_dim // 2]          # [..., d//2]
    x2 = x[..., head_dim // 2:]           # [..., d//2]
    rotated = torch.cat([-x2, x1], dim=-1)  # [..., d]
    return x * cos + rotated * sin


# ════════════════════════════════════════════════════════════
# 4. GQA 注意力（分组查询注意力）
# ════════════════════════════════════════════════════════════

class GroupedQueryAttention(nn.Module):
    """分组查询注意力 - 小模型=MHA，大模型节省KV Cache"""
    
    def __init__(self, config: XiaLiConfig):
        super().__init__()
        self.n_heads = config.n_heads
        self.n_kv_heads = config.n_kv_heads  # GQA: KV头数
        self.head_dim = config.d_model // config.n_heads
        self.n_rep = self.n_heads // self.n_kv_heads  # 每个KV头被几个Q头共享
        
        # 投影（无bias）
        self.wq = nn.Linear(config.d_model, config.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(config.d_model, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(config.d_model, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(config.n_heads * self.head_dim, config.d_model, bias=False)
        self.attn_dropout = nn.Dropout(config.dropout)
        
    def forward(
        self, 
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        batch, seq_len, _ = x.shape
        
        # QKV 投影
        q = self.wq(x).view(batch, seq_len, self.n_heads, self.head_dim)
        k = self.wk(x).view(batch, seq_len, self.n_kv_heads, self.head_dim)
        v = self.wv(x).view(batch, seq_len, self.n_kv_heads, self.head_dim)
        
        # RoPE
        cos_ = cos[:, :seq_len]
        sin_ = sin[:, :seq_len]
        q = apply_rope(q, cos_, sin_)
        k = apply_rope(k, cos_, sin_)
        
        # GQA: KV 头扩展到 Q 头数
        if self.n_rep > 1:
            k = k.repeat_interleave(self.n_rep, dim=2)
            v = v.repeat_interleave(self.n_rep, dim=2)
        
        # 转置为 [batch, n_heads, seq, head_dim]
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        
        # 注意力计算
        scale = self.head_dim ** -0.5
        attn = torch.matmul(q, k.transpose(-2, -1)) * scale
        
        if mask is not None:
            attn = attn + mask[:, :, :seq_len, :seq_len]
        
        attn = F.softmax(attn, dim=-1, dtype=torch.float32).to(q.dtype)
        attn = self.attn_dropout(attn)
        
        out = torch.matmul(attn, v)  # [batch, n_heads, seq, head_dim]
        out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1)
        
        return self.wo(out)


# ════════════════════════════════════════════════════════════
# 5. SwiGLU 前馈网络
# ════════════════════════════════════════════════════════════

class SwiGLU(nn.Module):
    """SwiGLU 前馈网络 - 比 ReLU FFN 质量更高"""
    
    def __init__(self, config: XiaLiConfig):
        super().__init__()
        # d_ff 已按 8/3*d_model 计算。SwiGLU: w1和w3同维度，w2降回
        hidden = config.d_ff
        self.w1 = nn.Linear(config.d_model, hidden, bias=False)
        self.w2 = nn.Linear(hidden, config.d_model, bias=False)
        self.w3 = nn.Linear(config.d_model, hidden, bias=False)
        self.dropout = nn.Dropout(config.dropout)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU(x) = (SiLU(x @ W1) * (x @ W3)) @ W2
        return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))


# ════════════════════════════════════════════════════════════
# 6. Transformer 块
# ════════════════════════════════════════════════════════════

class XiaLiBlock(nn.Module):
    """Transformer 块 - Pre-RMSNorm + GQA + SwiGLU"""
    
    def __init__(self, config: XiaLiConfig, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        
        self.attention = GroupedQueryAttention(config)
        self.feed_forward = SwiGLU(config)
        
        self.attn_norm = XiaLiRMSNorm(config.d_model, config.norm_eps)
        self.ff_norm = XiaLiRMSNorm(config.d_model, config.norm_eps)
        
    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        h = x + self.attention(self.attn_norm(x), cos, sin, mask)
        out = h + self.feed_forward(self.ff_norm(h))
        return out


# ════════════════════════════════════════════════════════════
# 7. 主模型
# ════════════════════════════════════════════════════════════

class XiaLiModel(nn.Module):
    """
    XiaLi-Next 统一架构
    
    核心设计：
    - RoPE 位置编码（无学习的 position embedding）
    - Pre-RMSNorm（训练稳定）
    - SwiGLU FFN（同等参数量更高质量）
    - GQA 注意力（大模型省 KV Cache）
    - Weight Tying（embedding 与输出层共享权重）
    - 所有 Linear 无 bias（更干净）
    - 从 20M 到 7B+ 统一架构
    """
    
    def __init__(self, config: XiaLiConfig):
        super().__init__()
        self.config = config
        
        # 词嵌入
        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        
        # 预计算 RoPE
        cos, sin = precompute_rope_freqs(
            config.d_model // config.n_heads,
            config.max_seq_len * 2,
            config.rope_theta,
        )
        self.register_buffer('rope_cos', cos, persistent=False)
        self.register_buffer('rope_sin', sin, persistent=False)
        
        # Transformer 块
        self.blocks = nn.ModuleList([
            XiaLiBlock(config, i) for i in range(config.n_layers)
        ])
        
        # 最终归一化
        self.norm = XiaLiRMSNorm(config.d_model, config.norm_eps)
        
        # 输出层（与 embedding 共享权重）
        self.fc_out = nn.Linear(config.d_model, config.vocab_size, bias=False)
        
        if config.tie_weights:
            self.fc_out.weight = self.embedding.weight  # Weight Tying
        
        self._init_weights()
    
    def _init_weights(self):
        """初始化权重"""
        for name, param in self.named_parameters():
            if param.dim() >= 2:
                nn.init.normal_(param, mean=0.0, std=0.02)
            else:
                nn.init.zeros_(param)
    
    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        batch, seq_len = x.shape
        
        # 词嵌入
        h = self.embedding(x)
        
        # RoPE
        cos = self.rope_cos[:, :seq_len]
        sin = self.rope_sin[:, :seq_len]
        
        # Transformer 层
        for block in self.blocks:
            h = block(h, cos, sin, mask)
        
        # 最终归一化
        h = self.norm(h)
        
        # 输出
        return self.fc_out(h)
    
    def generate_mask(self, seq_len: int) -> torch.Tensor:
        """生成因果掩码 [1, 1, seq, seq]"""
        mask = torch.full((seq_len, seq_len), float('-inf'))
        mask = torch.triu(mask, diagonal=1)
        return mask.unsqueeze(0).unsqueeze(0)
    
    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 100,
        temperature: float = 0.7,
        top_k: int = 40,
        eos_id: int = 3,
    ) -> torch.Tensor:
        """自回归生成"""
        self.eval()
        for _ in range(max_new_tokens):
            # 截断到 max_seq_len
            seq = input_ids if input_ids.size(1) <= self.config.max_seq_len else input_ids[:, -self.config.max_seq_len:]
            
            mask = self.generate_mask(seq.size(1)).to(seq.device)
            logits = self.forward(seq, mask)
            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')
            
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, 1)
            input_ids = torch.cat([input_ids, next_token], dim=-1)
            
            if next_token.item() == eos_id:
                break
        
        return input_ids
    
    def count_params(self) -> dict:
        """统计参数量"""
        total = sum(p.numel() for p in self.parameters())
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        # 各模块
        emb = self.embedding.weight.numel()
        attn = sum(p.numel() for n, p in self.named_parameters() if 'attention' in n)
        ffn = sum(p.numel() for n, p in self.named_parameters() if 'feed_forward' in n)
        norm = sum(p.numel() for n, p in self.named_parameters() if 'norm' in n)
        return {
            'total': total,
            'trainable': trainable,
            'embedding': emb,
            'attention': attn,
            'ffn': ffn,
            'norm': norm,
            'total_m': total / 1e6,
        }


# ════════════════════════════════════════════════════════════
# 8. 与原有系统兼容的训练适配
# ════════════════════════════════════════════════════════════

# 导出兼容名（方便从原系统切换）
XiaLiNextModel = XiaLiModel
XiaLiNextConfig = XiaLiConfig


# ════════════════════════════════════════════════════════════
# 9. 快速测试
# ════════════════════════════════════════════════════════════

if __name__ == "__main__":
    print("=" * 65)
    print("  XiaLi-Next 架构测试")
    print("=" * 65)
    
    for name in ['20m', '26m', '100m']:
        cfg = XiaLiConfig.preset(name)
        try:
            model = XiaLiModel(cfg)
            stats = model.count_params()
            print(f"\n  [{name}] {cfg.d_model}d × {cfg.n_layers}L × {cfg.n_heads}H (KV={cfg.n_kv_heads})")
            print(f"    参数量: {stats['total_m']:.2f}M")
            print(f"    其中: emb={stats['embedding']/1e6:.2f}M, attn={stats['attention']/1e6:.2f}M, "
                  f"ffn={stats['ffn']/1e6:.2f}M")
            
            # 前向测试
            x = torch.randint(0, cfg.vocab_size - 1, (1, 64))
            try:
                out = model(x)
                print(f"    前向: ✅ {out.shape}")
            except Exception as e:
                print(f"    前向: ❌ {e}")
        except Exception as e:
            print(f"\n  [{name}] ❌ {e}")
    
    print(f"\n{'='*65}")
    print(f"  ✅ XiaLi-Next 就绪！")
    print(f"  从 20M 到 7B+ 统一架构")
    print(f"{'='*65}")