/// Sinusoidal positional encoding (fixed, not learned).
///
/// Python:
///   position = torch.arange(max_seq_len).unsqueeze(1)
///   div_term = torch.exp(torch.arange(0, d_model, 2) * (-log(10000.0) / d_model))
///   pe[:, 0::2] = sin(position * div_term)
///   pe[:, 1::2] = cos(position * div_term)
///   pe = pe.unsqueeze(0)  # [1, max_seq_len, d_model]

use burn::prelude::*;

/// Generate sinusoidal positional encoding buffer.
///
/// Returns: [1, max_seq_len, d_model]
pub fn positional_encoding<B: Backend>(
    max_seq_len: usize,
    d_model: usize,
    device: &B::Device,
) -> Tensor<B, 3> {
    let half_d = d_model / 2;
    let mut pe = vec![0.0f32; max_seq_len * d_model];

    for pos in 0..max_seq_len {
        for i in 0..half_d {
            let div_term = (-(i as f64 * 2.0) * (10000.0_f64.ln()) / d_model as f64).exp() as f32;
            let angle = pos as f32 * div_term;
            pe[pos * d_model + i * 2] = angle.sin();
            pe[pos * d_model + i * 2 + 1] = angle.cos();
        }
    }

    Tensor::<B, 3>::from_data(
        TensorData::new(pe, vec![1, max_seq_len, d_model]),
        device,
    )
}

/// Add positional encoding to input tensor.
///
/// x: [B, S, D] — adds pe[:, :S, :] to x
pub fn add_positional_encoding<B: Backend>(
    x: Tensor<B, 3>,
    pe: &Tensor<B, 3>,
) -> Tensor<B, 3> {
    let seq_len = x.dims()[1];
    let pe_slice = pe.clone().slice([0..1, 0..seq_len]);
    x + pe_slice
}
