use burn::tensor::backend::Backend;
use burn::tensor::{Float, Tensor};

#[derive(Debug, Clone)]
pub struct RoPE {
    pub dim: usize,
    pub theta: f64,
    pub max_seq_len: usize,
    cos_cache: Vec<f32>,
    sin_cache: Vec<f32>,
    half_dim: usize,
}

impl RoPE {
    pub fn new(dim: usize, theta: f64, max_seq_len: usize) -> Self {
        let half_dim = dim / 2;

        let mut inv_freqs = Vec::with_capacity(half_dim);
        for i in 0..half_dim {
            let freq = 1.0 / (theta.powf((2.0 * i as f64) / dim as f64));
            inv_freqs.push(freq as f32);
        }

        let mut cos_cache = Vec::with_capacity(max_seq_len * half_dim);
        let mut sin_cache = Vec::with_capacity(max_seq_len * half_dim);

        for pos in 0..max_seq_len {
            for &freq in &inv_freqs {
                let angle = pos as f32 * freq;
                cos_cache.push(angle.cos());
                sin_cache.push(angle.sin());
            }
        }

        Self {
            dim,
            theta,
            max_seq_len,
            cos_cache,
            sin_cache,
            half_dim,
        }
    }

    pub fn get_cache(&self, offset: usize, len: usize) -> (Vec<f32>, Vec<f32>) {
        let cos = self.cos_cache[offset * self.half_dim..(offset + len) * self.half_dim].to_vec();
        let sin = self.sin_cache[offset * self.half_dim..(offset + len) * self.half_dim].to_vec();
        (cos, sin)
    }

    pub fn apply<B: Backend>(
        &self,
        x: Tensor<B, 4, Float>,
        offset: usize,
    ) -> Tensor<B, 4, Float> {
        let device = x.device();
        let shape = x.shape();
        let dims: [usize; 4] = shape.dims();
        let batch = dims[0];
        let heads = dims[1];
        let seq_len = dims[2];
        let dim = dims[3];
        let half_dim = self.half_dim;

        let cos_vals = &self.cos_cache[offset * half_dim..(offset + seq_len) * half_dim];
        let sin_vals = &self.sin_cache[offset * half_dim..(offset + seq_len) * half_dim];

        let cos = Tensor::<B, 1, Float>::from_floats(cos_vals, &device)
            .reshape([1, 1, seq_len, half_dim])
            .expand([batch, heads, seq_len, half_dim]);
        let sin = Tensor::<B, 1, Float>::from_floats(sin_vals, &device)
            .reshape([1, 1, seq_len, half_dim])
            .expand([batch, heads, seq_len, half_dim]);

        let x1 = x.clone().slice([0..batch, 0..heads, 0..seq_len, 0..half_dim]);
        let x2 = x.slice([0..batch, 0..heads, 0..seq_len, half_dim..dim]);

        let x1_rot = x1.clone() * cos.clone() - x2.clone() * sin.clone();
        let x2_rot = x2 * cos + x1 * sin;

        Tensor::cat(vec![x1_rot, x2_rot], 3)
    }
}