"""
位置编码模块 — Person B 负责实现

包含:
1. SinusoidalPositionalEncoding: 经典正弦余弦位置编码
2. RotaryPositionalEmbedding (RoPE): 旋转位置编码

RoPE 是当前最前沿的位置编码方案，被 LLaMA, GPT-NeoX 等大模型广泛采用。
优势: 更好地捕获相对位置信息，支持长度外推。
"""

from __future__ import annotations

import math

import torch
import torch.nn as nn


class SinusoidalPositionalEncoding(nn.Module):
    """
    经典正弦余弦位置编码 (Vaswani et al., 2017)。

    PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
    PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

    TODO [Person B]: 实现以下内容:
    1. 在 __init__ 中预计算位置编码矩阵 [max_seq_len, d_model]
    2. 注册为 buffer (不参与梯度更新)
    3. forward(x): 返回 x + pe[:, :x.size(1), :]
    """

    def __init__(self, d_model: int, max_seq_len: int = 5000, dropout: float = 0.1):
        super().__init__()
        raise NotImplementedError("TODO: Person B 实现 SinusoidalPositionalEncoding.__init__")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: [B, L, D]
        Returns:
            x + positional_encoding: [B, L, D]
        """
        raise NotImplementedError("TODO: Person B 实现 SinusoidalPositionalEncoding.forward")


class RotaryPositionalEmbedding(nn.Module):
    """
    旋转位置编码 (RoPE) — Su et al., 2021

    核心思想: 通过旋转变换将位置信息编码到 Q, K 向量中。
    q' = q * cos(θ) + rotate_half(q) * sin(θ)
    k' = k * cos(θ) + rotate_half(k) * sin(θ)

    TODO [Person B]: 实现以下内容:

    __init__:
    1. 计算频率: inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2) / dim))
    2. 注册为 buffer

    _compute_rope(seq_len):
    1. 计算 position indices: t = [0, 1, ..., seq_len-1]
    2. 计算 freqs = torch.outer(t, inv_freq)
    3. 构建 cos_cached, sin_cached

    apply_rotary_pos_emb(q, k):
    1. 对 q 和 k 应用旋转变换
    2. 返回旋转后的 q', k'

    参考: https://arxiv.org/abs/2104.09864
    """

    def __init__(self, dim: int, max_seq_len: int = 2048, base: float = 10000.0):
        super().__init__()
        raise NotImplementedError("TODO: Person B 实现 RotaryPositionalEmbedding.__init__")

    def _compute_rope(self, seq_len: int, device: torch.device):
        raise NotImplementedError("TODO: Person B 实现 _compute_rope")

    @staticmethod
    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
        """将 x 的后半部分取反并与前半部分交换。"""
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)

    def apply_rotary_pos_emb(
        self,
        q: torch.Tensor,  # [B, nhead, L, d_k]
        k: torch.Tensor,  # [B, nhead, L, d_k]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        对 Q, K 应用 RoPE。

        Returns:
            (q_rotated, k_rotated)
        """
        raise NotImplementedError("TODO: Person B 实现 apply_rotary_pos_emb")