//! Rotary Position Embedding (RoPE).

/// Precompute RoPE frequency table.
///
/// Returns (cos, sin) vectors of length [max_seq_len * head_dim/2].
pub fn precompute_freqs(head_dim: usize, max_seq_len: usize, base: f32) -> (Vec<f32>, Vec<f32>) {
    let half_dim = head_dim / 2;
    let mut cos_table = vec![0.0f32; max_seq_len * half_dim];
    let mut sin_table = vec![0.0f32; max_seq_len * half_dim];

    let base_f64 = base as f64;
    let head_dim_f64 = head_dim as f64;

    for pos in 0..max_seq_len {
        for i in 0..half_dim {
            // Use f64 for the frequency computation to avoid precision loss
            // with large rope_theta values (e.g. Qwen2.5 uses 1,000,000)
            let freq = 1.0_f64 / base_f64.powf(2.0 * i as f64 / head_dim_f64);
            let angle = pos as f64 * freq;
            cos_table[pos * half_dim + i] = angle.cos() as f32;
            sin_table[pos * half_dim + i] = angle.sin() as f32;
        }
    }

    (cos_table, sin_table)
}

/// Apply RoPE to Q or K tensor in-place.
///
/// `x` layout: [seq_len, num_heads, head_dim] flattened as [seq_len * num_heads * head_dim]
/// Applies rotation to pairs of elements in the head dimension.
pub fn apply_rope(
    x: &mut [f32],
    seq_len: usize,
    num_heads: usize,
    head_dim: usize,
    cos_table: &[f32],
    sin_table: &[f32],
    start_pos: usize,
) {
    let half_dim = head_dim / 2;
    for s in 0..seq_len {
        let pos = start_pos + s;
        for h in 0..num_heads {
            let offset = (s * num_heads + h) * head_dim;
            for i in 0..half_dim {
                let cos_val = cos_table[pos * half_dim + i];
                let sin_val = sin_table[pos * half_dim + i];
                let x0 = x[offset + i];
                let x1 = x[offset + half_dim + i];
                x[offset + i] = x0 * cos_val - x1 * sin_val;
                x[offset + half_dim + i] = x0 * sin_val + x1 * cos_val;
            }
        }
    }
}

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

    #[test]
    fn test_precompute_freqs_shape() {
        let (cos, sin) = precompute_freqs(8, 16, 10000.0);
        assert_eq!(cos.len(), 16 * 4); // max_seq_len * half_dim
        assert_eq!(sin.len(), 16 * 4);
    }

    #[test]
    fn test_rope_position_zero_is_identity_for_cos() {
        let (cos, _sin) = precompute_freqs(8, 16, 10000.0);
        // At position 0, cos should be 1.0, sin should be 0.0
        for (i, c) in cos.iter().enumerate().take(4) {
            assert!((c - 1.0).abs() < 1e-6, "cos[{i}] = {c}");
        }
    }

    #[test]
    fn test_rope_frequency_precision_high_base() {
        // Qwen2.5 uses rope_theta = 1_000_000.0
        let head_dim = 64;
        let max_seq_len = 8192;
        let base = 1_000_000.0_f32;

        let (cos_table, sin_table) = precompute_freqs(head_dim, max_seq_len, base);

        // At position 8191, the lowest frequency should still be distinguishable
        let half_dim = head_dim / 2;
        let last_pos = max_seq_len - 1;
        let cos_val = cos_table[last_pos * half_dim];
        let sin_val = sin_table[last_pos * half_dim];

        // cos² + sin² must equal 1.0 (rotation property)
        let sum = cos_val * cos_val + sin_val * sin_val;
        assert!((sum - 1.0).abs() < 1e-5, "Rotation property violated: cos²+sin²={sum}");

        // Adjacent positions must produce different values
        let prev_cos = cos_table[(last_pos - 1) * half_dim];
        assert!((cos_val - prev_cos).abs() > 1e-10, "Adjacent positions indistinguishable");
    }

    #[test]
    fn test_apply_rope_preserves_norm() {
        let head_dim = 8;
        let (cos, sin) = precompute_freqs(head_dim, 16, 10000.0);
        let mut x = vec![1.0f32; head_dim];
        let norm_before: f32 = x.iter().map(|v| v * v).sum::<f32>().sqrt();
        apply_rope(&mut x, 1, 1, head_dim, &cos, &sin, 5);
        let norm_after: f32 = x.iter().map(|v| v * v).sum::<f32>().sqrt();
        assert!(
            (norm_before - norm_after).abs() < 1e-4,
            "norm changed: {norm_before} -> {norm_after}"
        );
    }
}
