"""基础算子：RMSNorm、解耦 RoPE、SwiGLU FFN。"""

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


class RMSNorm(nn.Module):
    """比 LayerNorm 少一次均值统计，省算力且效果相当，是现代 LLM 的标配。

        y = x / sqrt(mean(x^2) + eps) * g
    """

    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:
        # 统计量始终用 fp32 算，避免半精度下 rsqrt 抖动
        out = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
        return (out * self.weight.float()).type_as(x)


# --------------------------------------------------------------------------
#  RoPE：旋转位置编码
#  MLA 里它是「解耦」的 —— 只作用在 Q/K 单独切出来的 qk_rope_head_dim 维上，
#  剩下的 qk_nope 部分不带位置信息，这样 K 才能被压进和位置无关的潜在向量里。
# --------------------------------------------------------------------------

def build_rope_cache(head_dim: int, max_seq_len: int, theta: float = 10000.0,
                     scaling: float = 1.0, device=None):
    """返回 (cos, sin)，形状均为 (max_seq_len, head_dim)。

    scaling > 1 时做 NTK-aware 频率插值：把 base 放大 scaling^(d/(d-2))，
    等价于低频维度插值、高频维度基本不动，用来做上下文长度外推。
    """
    if scaling > 1.0:
        theta = theta * (scaling ** (head_dim / (head_dim - 2)))
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim))
    t = torch.arange(max_seq_len, dtype=torch.float32, device=device)
    freqs = torch.outer(t, inv_freq)                 # (T, head_dim/2)
    emb = torch.cat([freqs, freqs], dim=-1)          # (T, head_dim)
    return emb.cos(), emb.sin()


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """x: (B, H, T, D)，cos/sin: (T, D)。"""
    cos = cos[None, None, :, :].to(x.dtype)
    sin = sin[None, None, :, :].to(x.dtype)
    return x * cos + rotate_half(x) * sin


# --------------------------------------------------------------------------

class SwiGLU(nn.Module):
    """FFN(x) = W_down( silu(W_gate x) * W_up x )，比 ReLU-MLP 表达力更强。"""

    def __init__(self, dim: int, inter_dim: int):
        super().__init__()
        self.w_gate = nn.Linear(dim, inter_dim, bias=False)
        self.w_up = nn.Linear(dim, inter_dim, bias=False)
        self.w_down = nn.Linear(inter_dim, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))