//! Compute backend abstraction for LLMHosts Core.
//!
//! The `ComputeBackend` trait provides a uniform interface over CPU, CUDA,
//! ROCm (via WGPU), Metal (via WGPU), and Vulkan (via WGPU) backends.
//!
//! Default feature set compiles only the CPU backend (zero extra deps).
//! Enable `cuda` or `wgpu` features to unlock GPU acceleration.

pub mod auto;
pub mod cpu;

#[cfg(feature = "cuda")]
pub mod cuda;

#[cfg(feature = "wgpu")]
pub mod wgpu_backend;

use crate::tensor::Tensor;

/// Abstraction over hardware-accelerated compute backends.
///
/// Every backend must implement the core operations needed by the
/// LLMHosts transformer forward pass. The CPU backend delegates to
/// the existing hand-written kernels; GPU backends will use CubeCL.
pub trait ComputeBackend: Send + Sync {
    /// Human-readable name for logging (e.g. "cpu", "cuda:0", "wgpu:Metal").
    fn name(&self) -> &str;

    /// Matrix multiply: A [M, K] x B [K, N] -> C [M, N].
    fn matmul(&self, a: &Tensor, b: &Tensor) -> Tensor;

    /// Softmax over the given dimension (in-place).
    fn softmax(&self, x: &mut Tensor, dim: usize);

    /// Apply Rotary Position Embedding to Q and K tensors.
    ///
    /// `q` layout: [seq_len, num_heads * head_dim]
    /// `k` layout: [seq_len, num_kv_heads * head_dim]
    /// `freqs_cos` / `freqs_sin`: precomputed RoPE tables.
    #[allow(clippy::too_many_arguments)]
    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,
    );

    /// RMS normalization with learned weight vector.
    ///
    /// Normalizes `x` by its RMS, then element-wise multiplies by `weight`.
    fn rmsnorm(&self, x: &Tensor, weight: &Tensor, eps: f32) -> Tensor;

    /// SiLU-gated element-wise multiply: gate = silu(gate) * up.
    ///
    /// `gate` is modified in-place.
    fn silu_elementwise_mul(&self, gate: &mut Tensor, up: &Tensor);
}

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

    /// Verify the CPU backend is returned by default auto-detect.
    #[test]
    fn test_auto_detect_returns_cpu_by_default() {
        let backend = auto::detect_best_backend();
        assert_eq!(backend.name(), "cpu");
    }

    /// Verify the ComputeBackend trait is object-safe (can be boxed).
    #[test]
    fn test_trait_is_object_safe() {
        let backend: Box<dyn ComputeBackend> = Box::new(cpu::CpuBackend::new());
        assert_eq!(backend.name(), "cpu");
    }

    /// Round-trip: matmul through the backend matches direct tensor::matmul_2d.
    #[test]
    fn test_cpu_backend_matmul_matches_direct() {
        let backend = cpu::CpuBackend::new();
        let a = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
        let b = Tensor::from_data(vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0], vec![3, 2]).unwrap();

        let via_backend = backend.matmul(&a, &b);
        let via_direct = crate::tensor::matmul_2d(&a, &b).unwrap();

        assert_eq!(via_backend.shape, via_direct.shape);
        for (a, b) in via_backend.data.iter().zip(via_direct.data.iter()) {
            assert!((a - b).abs() < 1e-5, "mismatch: {a} vs {b}");
        }
    }

    /// Softmax via backend produces a valid probability distribution.
    #[test]
    fn test_cpu_backend_softmax() {
        let backend = cpu::CpuBackend::new();
        let mut t = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
        let last_dim = t.shape.len() - 1;
        backend.softmax(&mut t, last_dim);

        // Each row should sum to 1.0
        for row in t.data.chunks(3) {
            let sum: f32 = row.iter().sum();
            assert!((sum - 1.0).abs() < 1e-5, "row sum = {sum}");
        }
    }

    /// RMSNorm via backend produces finite, non-zero output.
    #[test]
    fn test_cpu_backend_rmsnorm() {
        let backend = cpu::CpuBackend::new();
        let x = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], vec![1, 4]).unwrap();
        let weight = Tensor::from_data(vec![1.0, 1.0, 1.0, 1.0], vec![4]).unwrap();
        let out = backend.rmsnorm(&x, &weight, 1e-6);
        assert_eq!(out.shape, x.shape);
        assert!(out.data.iter().all(|v| v.is_finite()));
        assert!(out.data.iter().any(|v| *v != 0.0));
    }

    /// SiLU element-wise mul via backend.
    #[test]
    fn test_cpu_backend_silu_elementwise_mul() {
        let backend = cpu::CpuBackend::new();
        let mut gate = Tensor::from_data(vec![1.0, 2.0, -1.0], vec![3]).unwrap();
        let up = Tensor::from_data(vec![1.0, 1.0, 1.0], vec![3]).unwrap();
        backend.silu_elementwise_mul(&mut gate, &up);
        // silu(1)*1 ≈ 0.731, silu(2)*1 ≈ 1.762, silu(-1)*1 ≈ -0.269
        assert!((gate.data[0] - 0.731).abs() < 0.01);
        assert!((gate.data[1] - 1.762).abs() < 0.01);
        assert!((gate.data[2] - (-0.269)).abs() < 0.01);
    }

    /// RoPE via backend preserves vector norm.
    #[test]
    fn test_cpu_backend_rope_preserves_norm() {
        let backend = cpu::CpuBackend::new();
        let head_dim = 8;
        let num_heads = 2;
        let seq_len = 1;
        let data_len = seq_len * num_heads * head_dim;

        let mut q = Tensor::from_data(vec![1.0; data_len], vec![seq_len, num_heads * head_dim]).unwrap();
        let mut k = Tensor::from_data(vec![1.0; data_len], vec![seq_len, num_heads * head_dim]).unwrap();

        let norm_before: f32 = q.data.iter().map(|v| v * v).sum::<f32>().sqrt();

        let (cos_table, sin_table) = crate::rope::precompute_freqs(head_dim, 16, 10000.0);
        backend.rope(&mut q, &mut k, &cos_table, &sin_table, num_heads, num_heads, head_dim, 3);

        let norm_after: f32 = q.data.iter().map(|v| v * v).sum::<f32>().sqrt();
        assert!(
            (norm_before - norm_after).abs() < 1e-4,
            "norm changed: {norm_before} -> {norm_after}"
        );
    }
}
