//! CubeCL Rotary Position Embedding (RoPE) kernel.
//!
//! Applies rotary positional embeddings to Q or K tensors. Each work-unit
//! processes one (position, head, half-pair) triple, rotating a pair of
//! elements by the precomputed cos/sin angle.
//!
//! Generic over `cubecl::prelude::Float`, so the same kernel compiles for
//! f32, f16, bf16 on any CubeCL-supported backend (CUDA, WGPU, etc.).

use cubecl::prelude::*;

/// RoPE kernel: one work-unit per (position, head, half-dim-pair).
///
/// Layout (row-major, contiguous):
///   x:         [seq_len, num_heads * head_dim] — modified in-place via output
///   cos_table: [max_seq_len * half_dim]
///   sin_table: [max_seq_len * half_dim]
///   output:    [seq_len, num_heads * head_dim] — write-back buffer
///
/// Launch grid: (seq_len * num_heads * half_dim) work-units.
///
/// `num_heads`, `head_dim`, `half_dim`, and `start_pos` are `#[comptime]`.
#[cube(launch)]
pub fn rope_kernel<F: Float>(
    x: &Tensor<F>,
    cos_table: &Tensor<F>,
    sin_table: &Tensor<F>,
    output: &mut Tensor<F>,
    #[comptime] num_heads: u32,
    #[comptime] head_dim: u32,
    #[comptime] half_dim: u32,
    #[comptime] start_pos: u32,
) {
    let idx = ABSOLUTE_POS;
    let total_per_pos = num_heads * half_dim;

    // Decompose flat index into (s, h, i)
    let s = idx / total_per_pos;
    let remainder = idx % total_per_pos;
    let h = remainder / half_dim;
    let i = remainder % half_dim;

    let pos = start_pos + s;
    let offset = (s * num_heads + h) * head_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];

    output[offset + i] = x0 * cos_val - x1 * sin_val;
    output[offset + half_dim + i] = x0 * sin_val + x1 * cos_val;
}

/// Launch configuration for the RoPE kernel.
///
/// Returns `(cube_count, cube_dim)` suitable for passing to the generated
/// `rope_kernel::launch` function.
///
/// One thread per (position, head, half-dim-pair), launched in 1D blocks of 256.
pub fn rope_launch_config(seq_len: u32, num_heads: u32, half_dim: u32) -> (CubeCount, CubeDim) {
    let total = seq_len * num_heads * half_dim;
    let block_size = 256u32;
    let cube_dim = CubeDim::new(block_size, 1, 1);
    let cube_count = CubeCount::Static((total + block_size - 1) / block_size, 1, 1);
    (cube_count, cube_dim)
}

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

    fn cube_count_xyz(c: CubeCount) -> (u32, u32, u32) {
        match c {
            CubeCount::Static(x, y, z) => (x, y, z),
            _ => panic!("expected CubeCount::Static"),
        }
    }

    #[test]
    fn test_rope_launch_config_basic() {
        // seq_len=1, num_heads=8, half_dim=32 → total=256
        let (count, _dim) = rope_launch_config(1, 8, 32);
        assert_eq!(cube_count_xyz(count), (1, 1, 1));
    }

    #[test]
    fn test_rope_launch_config_multiple_positions() {
        // seq_len=4, num_heads=8, half_dim=32 → total=1024
        let (count, _dim) = rope_launch_config(4, 8, 32);
        assert_eq!(cube_count_xyz(count), (4, 1, 1));
    }

    #[test]
    fn test_rope_launch_config_single_head() {
        // seq_len=1, num_heads=1, half_dim=4 → total=4
        let (count, _dim) = rope_launch_config(1, 1, 4);
        assert_eq!(cube_count_xyz(count), (1, 1, 1));
    }

    #[test]
    fn test_rope_launch_config_non_aligned() {
        // seq_len=3, num_heads=2, half_dim=4 → total=24
        let (count, _dim) = rope_launch_config(3, 2, 4);
        assert_eq!(cube_count_xyz(count), (1, 1, 1));
    }
}
