//! 旋转位置编码 (Rotary Position Embedding, RoPE)
//!
//! RoPE 通过旋转变换将位置信息编码到注意力机制中。
//! 这种方法具有良好的外推能力，是 Gemma/Llama 系列模型的标准位置编码方式。
//!
//! # 数学原理
//!
//! 对于位置 `m` 和维度索引 `i`，RoPE 将向量旋转角度 `θ_i * m`，其中：
//! - `θ_i = 10000^(-2i/d)`
//!
//! 旋转公式：
//! ```text
//! [x1', x2'] = [x1*cos(θ) - x2*sin(θ), x1*sin(θ) + x2*cos(θ)]
//! ```

use candle_core::{DType, Device, Result, Tensor, D};

/// 旋转位置编码
pub struct RotaryEmbedding {
    /// 预计算的 cos 值: (max_seq_len, head_dim)
    cos: Tensor,
    /// 预计算的 sin 值: (max_seq_len, head_dim)
    sin: Tensor,
    /// 每个注意力头的维度
    head_dim: usize,
    /// 数据类型
    dtype: DType,
}

impl RotaryEmbedding {
    /// 创建 RoPE 实例
    ///
    /// # Arguments
    /// * `head_dim` - 每个注意力头的维度
    /// * `max_seq_len` - 最大序列长度
    /// * `theta` - 基础频率参数 (通常为 10000.0)
    /// * `dtype` - 数据类型
    /// * `device` - 计算设备
    pub fn new(
        head_dim: usize,
        max_seq_len: usize,
        theta: f64,
        dtype: DType,
        device: &Device,
    ) -> Result<Self> {
        // 计算逆频率: inv_freq[i] = 1 / (theta^(2i/d))
        let half_dim = head_dim / 2;
        let inv_freq: Vec<f32> = (0..half_dim)
            .map(|i| 1.0 / (theta as f32).powf((2 * i) as f32 / head_dim as f32))
            .collect();

        let inv_freq = Tensor::from_vec(inv_freq, (half_dim,), device)?;

        // 位置序列: [0, 1, 2, ..., max_seq_len-1]
        let positions: Vec<f32> = (0..max_seq_len).map(|i| i as f32).collect();
        let positions = Tensor::from_vec(positions, (max_seq_len,), device)?;

        // 计算角度: angles[pos, i] = pos * inv_freq[i]
        // positions: (max_seq_len,) -> (max_seq_len, 1)
        // inv_freq: (half_dim,) -> (1, half_dim)
        // angles: (max_seq_len, half_dim)
        let angles = positions
            .unsqueeze(1)?
            .broadcast_mul(&inv_freq.unsqueeze(0)?)?;

        // 复制角度以匹配完整的 head_dim
        // (max_seq_len, half_dim) -> (max_seq_len, head_dim)
        let angles = Tensor::cat(&[&angles, &angles], 1)?;

        // 预计算 cos 和 sin
        let cos = angles.cos()?.to_dtype(dtype)?;
        let sin = angles.sin()?.to_dtype(dtype)?;

        Ok(Self {
            cos,
            sin,
            head_dim,
            dtype,
        })
    }

    /// 应用 RoPE 到 Q 和 K 张量
    ///
    /// # Arguments
    /// * `q` - Query 张量，形状: (batch, num_heads, seq_len, head_dim)
    /// * `k` - Key 张量，形状: (batch, num_heads, seq_len, head_dim)
    /// * `start_pos` - 起始位置 (用于增量解码，embedding 场景通常为 0)
    ///
    /// # Returns
    /// 旋转后的 (q, k) 元组
    pub fn apply(&self, q: &Tensor, k: &Tensor, start_pos: usize) -> Result<(Tensor, Tensor)> {
        let seq_len = q.dim(2)?;

        // 获取当前位置范围的 cos 和 sin
        let cos = self.cos.narrow(0, start_pos, seq_len)?;
        let sin = self.sin.narrow(0, start_pos, seq_len)?;

        // 应用旋转
        let q_rotated = self.apply_rotary_emb(q, &cos, &sin)?;
        let k_rotated = self.apply_rotary_emb(k, &cos, &sin)?;

        Ok((q_rotated, k_rotated))
    }

    /// 对单个张量应用旋转位置编码
    ///
    /// # Arguments
    /// * `x` - 输入张量，形状: (batch, num_heads, seq_len, head_dim)
    /// * `cos` - cos 值，形状: (seq_len, head_dim)
    /// * `sin` - sin 值，形状: (seq_len, head_dim)
    fn apply_rotary_emb(&self, x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result<Tensor> {
        let half_dim = self.head_dim / 2;

        // 将 x 分成两半
        // x: (batch, num_heads, seq_len, head_dim)
        let x1 = x.narrow(D::Minus1, 0, half_dim)?;
        let x2 = x.narrow(D::Minus1, half_dim, half_dim)?;

        // cos, sin: (seq_len, head_dim) -> (1, 1, seq_len, head_dim)
        let cos = cos.unsqueeze(0)?.unsqueeze(0)?;
        let sin = sin.unsqueeze(0)?.unsqueeze(0)?;

        // 分别获取 cos/sin 的前半部分
        let cos_half = cos.narrow(D::Minus1, 0, half_dim)?;
        let sin_half = sin.narrow(D::Minus1, 0, half_dim)?;

        // 旋转公式:
        // x1' = x1 * cos - x2 * sin
        // x2' = x1 * sin + x2 * cos
        let rotated_x1 = (x1.broadcast_mul(&cos_half)? - x2.broadcast_mul(&sin_half)?)?;
        let rotated_x2 = (x1.broadcast_mul(&sin_half)? + x2.broadcast_mul(&cos_half)?)?;

        // 拼接回完整维度
        Tensor::cat(&[rotated_x1, rotated_x2], D::Minus1)
    }

    /// 获取 head_dim
    pub fn head_dim(&self) -> usize {
        self.head_dim
    }

    /// 获取数据类型
    pub fn dtype(&self) -> DType {
        self.dtype
    }
}

/// 简化的 RoPE 应用函数（用于测试）
#[allow(dead_code)]
pub fn apply_rope_simple(
    x: &Tensor,
    cos: &Tensor,
    sin: &Tensor,
    head_dim: usize,
) -> Result<Tensor> {
    let half_dim = head_dim / 2;

    let x1 = x.narrow(D::Minus1, 0, half_dim)?;
    let x2 = x.narrow(D::Minus1, half_dim, half_dim)?;

    let cos_half = cos.narrow(D::Minus1, 0, half_dim)?;
    let sin_half = sin.narrow(D::Minus1, 0, half_dim)?;

    let rotated_x1 = (x1.broadcast_mul(&cos_half)? - x2.broadcast_mul(&sin_half)?)?;
    let rotated_x2 = (x1.broadcast_mul(&sin_half)? + x2.broadcast_mul(&cos_half)?)?;

    Tensor::cat(&[rotated_x1, rotated_x2], D::Minus1)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rope_creation() {
        let device = Device::Cpu;
        let rope = RotaryEmbedding::new(64, 128, 10000.0, DType::F32, &device);
        assert!(rope.is_ok());

        let rope = rope.unwrap();
        assert_eq!(rope.head_dim(), 64);
    }

    #[test]
    fn test_rope_cos_sin_shape() {
        let device = Device::Cpu;
        let rope = RotaryEmbedding::new(64, 128, 10000.0, DType::F32, &device).unwrap();

        // cos 和 sin 应该是 (max_seq_len, head_dim)
        assert_eq!(rope.cos.dims(), &[128, 64]);
        assert_eq!(rope.sin.dims(), &[128, 64]);
    }

    #[test]
    fn test_rope_apply() {
        let device = Device::Cpu;
        let rope = RotaryEmbedding::new(64, 128, 10000.0, DType::F32, &device).unwrap();

        // 创建测试张量 (batch=2, heads=4, seq=16, head_dim=64)
        let q = Tensor::randn(0.0f32, 1.0, (2, 4, 16, 64), &device).unwrap();
        let k = Tensor::randn(0.0f32, 1.0, (2, 4, 16, 64), &device).unwrap();

        let result = rope.apply(&q, &k, 0);
        assert!(result.is_ok());

        let (q_rotated, k_rotated) = result.unwrap();
        assert_eq!(q_rotated.dims(), q.dims());
        assert_eq!(k_rotated.dims(), k.dims());
    }

    #[test]
    fn test_rope_preserves_norm() {
        // RoPE 应该保持向量的 L2 范数
        let device = Device::Cpu;
        let rope = RotaryEmbedding::new(64, 128, 10000.0, DType::F32, &device).unwrap();

        let q = Tensor::randn(0.0f32, 1.0, (1, 1, 1, 64), &device).unwrap();
        let k = q.clone();

        let (q_rotated, _) = rope.apply(&q, &k, 0).unwrap();

        // 计算原始和旋转后的范数
        let original_norm: f32 = q.sqr().unwrap().sum_all().unwrap().to_scalar().unwrap();
        let rotated_norm: f32 = q_rotated.sqr().unwrap().sum_all().unwrap().to_scalar().unwrap();

        // 范数应该相等（允许小误差）
        assert!((original_norm - rotated_norm).abs() < 1e-5);
    }
}
