//! CPU compute backend — wraps existing hand-written kernels.
//!
//! This backend compiles with zero optional dependencies (default features).
//! It delegates to the functions in `tensor`, `rope`, and `ffn` modules.

use crate::compute::ComputeBackend;
use crate::rope;
use crate::tensor::{matmul_2d, silu_inplace, Tensor};

/// CPU-only compute backend using matrixmultiply for BLAS-level matmul.
pub struct CpuBackend;

impl CpuBackend {
    pub fn new() -> Self {
        Self
    }
}

impl Default for CpuBackend {
    fn default() -> Self {
        Self::new()
    }
}

impl ComputeBackend for CpuBackend {
    fn name(&self) -> &str {
        "cpu"
    }

    fn matmul(&self, a: &Tensor, b: &Tensor) -> Tensor {
        matmul_2d(a, b).expect("matmul_2d shape mismatch")
    }

    fn softmax(&self, x: &mut Tensor, _dim: usize) {
        // Current implementation always operates on the last dimension.
        // The `dim` parameter is accepted for API compatibility with GPU
        // backends that can softmax over arbitrary axes.
        x.softmax_last_dim();
    }

    fn rope(
        &self,
        q: &mut Tensor,
        k: &mut Tensor,
        freqs_cos: &[f32],
        freqs_sin: &[f32],
        num_heads: usize,
        num_kv_heads: usize,
        head_dim: usize,
        start_pos: usize,
    ) {
        let seq_len = q.shape[0];
        rope::apply_rope(
            &mut q.data,
            seq_len,
            num_heads,
            head_dim,
            freqs_cos,
            freqs_sin,
            start_pos,
        );
        rope::apply_rope(
            &mut k.data,
            seq_len,
            num_kv_heads,
            head_dim,
            freqs_cos,
            freqs_sin,
            start_pos,
        );
    }

    fn rmsnorm(&self, x: &Tensor, weight: &Tensor, eps: f32) -> Tensor {
        let mut normed = x.rms_norm(eps);
        let hidden = weight.data.len();
        for chunk in normed.data.chunks_mut(hidden) {
            for (v, w) in chunk.iter_mut().zip(weight.data.iter()) {
                *v *= w;
            }
        }
        normed
    }

    fn silu_elementwise_mul(&self, gate: &mut Tensor, up: &Tensor) {
        silu_inplace(gate);
        gate.mul_inplace(up).expect("silu_elementwise_mul shape mismatch");
    }
}

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

    #[test]
    fn test_cpu_backend_name() {
        let b = CpuBackend::new();
        assert_eq!(b.name(), "cpu");
    }

    #[test]
    fn test_cpu_backend_default() {
        let b = CpuBackend;
        assert_eq!(b.name(), "cpu");
    }

    #[test]
    fn test_matmul_identity() {
        let b = CpuBackend::new();
        // Multiply by identity-ish: A [2,2] x I [2,2]
        let a = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
        let eye = Tensor::from_data(vec![1.0, 0.0, 0.0, 1.0], vec![2, 2]).unwrap();
        let c = b.matmul(&a, &eye);
        assert_eq!(c.shape, vec![2, 2]);
        assert!((c.data[0] - 1.0).abs() < 1e-6);
        assert!((c.data[1] - 2.0).abs() < 1e-6);
        assert!((c.data[2] - 3.0).abs() < 1e-6);
        assert!((c.data[3] - 4.0).abs() < 1e-6);
    }

    #[test]
    fn test_softmax_single_row() {
        let b = CpuBackend::new();
        let mut t = Tensor::from_data(vec![0.0, 0.0, 0.0], vec![1, 3]).unwrap();
        b.softmax(&mut t, 1);
        // Uniform input -> uniform output
        for v in &t.data {
            assert!((*v - 1.0 / 3.0).abs() < 1e-5);
        }
    }

    #[test]
    fn test_rmsnorm_unit_weight() {
        let b = CpuBackend::new();
        let x = Tensor::from_data(vec![3.0, 4.0], vec![1, 2]).unwrap();
        let w = Tensor::from_data(vec![1.0, 1.0], vec![2]).unwrap();
        let out = b.rmsnorm(&x, &w, 1e-6);
        // RMS of [3,4] = sqrt((9+16)/2) = sqrt(12.5) ≈ 3.536
        // normed ≈ [3/3.536, 4/3.536] ≈ [0.849, 1.131]
        assert!((out.data[0] - 0.849).abs() < 0.01, "got {}", out.data[0]);
        assert!((out.data[1] - 1.131).abs() < 0.01, "got {}", out.data[1]);
    }

    #[test]
    fn test_silu_elementwise_mul_zeros() {
        let b = CpuBackend::new();
        let mut gate = Tensor::from_data(vec![0.0, 0.0], vec![2]).unwrap();
        let up = Tensor::from_data(vec![5.0, 10.0], vec![2]).unwrap();
        b.silu_elementwise_mul(&mut gate, &up);
        // silu(0) = 0, so result = 0
        assert!((gate.data[0]).abs() < 1e-6);
        assert!((gate.data[1]).abs() < 1e-6);
    }
}
