#位置编码 (Positional Encoding)：给没有顺序的矩阵注入"位置感"

import torch
import torch.nn as nn
import math

class SinCosPositionalEncoding(nn.Module):
    """原始 Transformer 论文中的正弦/余弦绝对位置编码"""
    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
        """
        Args:
            d_model: Transformer 输入/输出的固定向量维度（论文里用了 512）
            max_len: 预定义的最大序列长度（超过这个长度需要额外扩展，但原始论文的性质允许一定的泛化）
            dropout: 防止过拟合的 Dropout 比例
        """
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        # 第一步：初始化一个全 0 的位置编码矩阵，形状是 (max_len, d_model)
        position_enc = torch.zeros(max_len, d_model)

        # 第二步：生成「位置序列」，把每个位置转换成浮点数后，多扩展一个维度（方便后面广播计算）
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)

        # 第三步：生成「频率衰减因子」——高频信号对应向量的低维度，低频信号对应高维度
        # 这样可以让不同位置的向量在高维度和低维度上都有差异，覆盖不同范围的位置关系
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )

        # 第四步：分别填充正弦和余弦部分——偶数维度用 sin，奇数维度用 cos
        position_enc[:, 0::2] = torch.sin(position * div_term)
        position_enc[:, 1::2] = torch.cos(position * div_term)

        # 第五步：把位置编码矩阵扩展成 (1, max_len, d_model) 的形状
        # 这样可以直接和形状为 (batch_size, seq_len, d_model) 的输入词向量相加（PyTorch 的广播机制）
        position_enc = position_enc.unsqueeze(0)

        # 第六步：把位置编码存进 buffer（训练时不会被优化，但会保存到模型文件里）
        self.register_buffer("position_enc", position_enc)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 输入的词向量矩阵，形状是 (batch_size, seq_len, d_model)
        Returns:
            加上位置编码后的最终表示向量，形状和输入完全一致
        """
        # 只取前 seq_len 个位置的编码（防止输入序列比预定义的 max_len 短）
        x = x + self.position_enc[:, :x.size(1), :]
        return self.dropout(x)

import torch
import torch.nn as nn

class LearnedPositionalEncoding(nn.Module):
    """端到端可学习的绝对位置编码"""
    def __init__(self, d_model: int, max_len: int = 5000):
        """
        Args:
            d_model: Transformer 输入/输出的固定向量维度
            max_len: 预定义的最大序列长度（**推理时不能超过这个长度！** 这点不如正弦余弦灵活）
        """
        super().__init__()
        # 用 Embedding 层存储位置向量：输入是位置索引（0 到 max_len-1），输出是 d_model 维的位置向量
        self.position_emb = nn.Embedding(max_len, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 输入的词向量矩阵，形状是 (batch_size, seq_len, d_model)
        Returns:
            加上位置编码后的最终表示向量，形状和输入完全一致
        """
        batch_size, seq_len, _ = x.shape

        # 生成当前 batch 需要的位置索引：0 到 seq_len-1，注意要和输入在同一个设备上（CPU/GPU/TPU）
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).repeat(batch_size, 1)

        # 取出位置向量，和词向量相加
        position_vectors = self.position_emb(positions)
        return x + position_vectors

import torch
import torch.nn as nn

class RotaryPositionalEmbedding(nn.Module):
    """现代 LLM 标配的 RoPE 旋转位置编码（核心简化版）"""
    def __init__(self, head_dim: int, base: int = 10000):
        """
        Args:
            head_dim: 多头注意力中每个头的向量维度（比如 d_model=512，num_heads=8，那么 head_dim=64）
            base: 频率衰减的基数（论文里用了 10000，LLaMA 系列也沿用了这个值）
        """
        super().__init__()
        # 预计算频率衰减因子，和正弦余弦编码里的 div_term 逻辑一样
        self.inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))

    def forward(self, x: torch.Tensor, seq_len: int = None) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            x: 输入的 Q/K 向量，形状是 (batch_size, num_heads, seq_len, head_dim)
            seq_len: 当前序列的长度（如果 x 已经是完整的 Q/K，也可以直接取 x.size(2)）
        Returns:
            cos_emb, sin_emb: 旋转用的余弦和正弦向量，形状是 (seq_len, head_dim)
        """
        if seq_len is None:
            seq_len = x.size(2)

        # 生成位置序列
        t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)

        # 计算每个位置、每个维度组的旋转角度
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)

        # 把正弦和余弦向量拼接起来（对应两两分组的 Q/K）
        cos_emb = torch.cat([freqs.cos(), freqs.cos()], dim=-1)
        sin_emb = torch.cat([freqs.sin(), freqs.sin()], dim=-1)

        return cos_emb, sin_emb