//! 计算后端抽象层 — SiFu 硬件加速
//!
//! Phase-51: 提供多层加速策略：
//! - **L0**: 手写 3×3 展开（零开销，所有平台）
//! - **L1**: SIMD 内联（≤32×32，桌面/服务端）
//! - **L2**: candle-core CPU 后端（大矩阵，需 `candle` feature）
//! - **L3**: candle CUDA / burn WGPU GPU 后端（需对应 feature）
//!
//! 所有后端实现统一的 `ComputeBackend` trait，`SiFuQueryEngine` 可透明切换。

use crate::query::sifu::{ActivationFn, TransformMatrix};

// ============================================================================
// ComputeBackend trait
// ============================================================================

/// 计算后端抽象
///
/// 封装 SiFu 信号传播所需的核心数值运算：
/// - 批量矩阵-向量乘法（信号转换）
/// - 批量激活函数
/// - 能量衰减
pub trait ComputeBackend: Send + Sync {
    /// 后端名称（用于日志和调试）
    fn name(&self) -> &str;

    /// 单个矩阵 × 向量乘法
    ///
    /// `matrix`: 行主序 (rows × cols) 矩阵
    /// `vector`: 长度为 cols 的输入向量
    /// `bias`:   可选偏置向量（长度为 rows），None 则无偏置
    ///
    /// 返回长度为 rows 的结果向量
    fn mat_mul_vec(
        &self,
        matrix: &TransformMatrix,
        vector: &[f32],
        bias: Option<&[f32]>,
    ) -> Vec<f32>;

    /// 批量矩阵-向量乘法
    ///
    /// 对 N 组 (matrix, vector, bias) 并行执行矩阵-向量乘法。
    /// 默认实现：逐个调用 `mat_mul_vec`。
    fn batch_mat_mul_vec(
        &self,
        matrices: &[&TransformMatrix],
        vectors: &[&[f32]],
        biases: &[Option<&[f32]>],
    ) -> Vec<Vec<f32>> {
        debug_assert_eq!(matrices.len(), vectors.len());
        debug_assert_eq!(matrices.len(), biases.len());
        matrices
            .iter()
            .zip(vectors.iter())
            .zip(biases.iter())
            .map(|((m, v), b)| self.mat_mul_vec(m, v, *b))
            .collect()
    }

    /// 批量应用激活函数（in-place）
    fn activate_batch(&self, values: &mut [f64], activation: ActivationFn) {
        for v in values.iter_mut() {
            *v = activation.apply(*v);
        }
    }

    /// 批量能量衰减 + 信号强度乘法
    ///
    /// `energies[i] = energies[i] * signal_strengths[i] * decay_factor`
    fn apply_decay(
        &self,
        energies: &mut [f64],
        signal_strengths: &[f64],
        decay_factor: f64,
    ) {
        debug_assert_eq!(energies.len(), signal_strengths.len());
        for (e, s) in energies.iter_mut().zip(signal_strengths.iter()) {
            *e *= s * decay_factor;
        }
    }

    /// 是否支持 GPU 加速
    fn is_gpu(&self) -> bool {
        false
    }

    /// 后端信息摘要
    fn info(&self) -> ComputeBackendInfo {
        ComputeBackendInfo {
            name: self.name().to_string(),
            is_gpu: self.is_gpu(),
            max_efficient_dim: self.max_efficient_dim(),
        }
    }

    /// 该后端高效处理的最大矩阵维度
    fn max_efficient_dim(&self) -> usize {
        usize::MAX
    }
}

/// 后端信息
#[derive(Debug, Clone)]
pub struct ComputeBackendInfo {
    /// 后端名称
    pub name: String,
    /// 是否 GPU 后端
    pub is_gpu: bool,
    /// 最大高效处理维度
    pub max_efficient_dim: usize,
}

// ============================================================================
// L0/L1: CpuBackend — 纯 Rust 实现
// ============================================================================

/// 纯 Rust CPU 计算后端
///
/// - 3×3 矩阵：手写展开（L0，零开销）
/// - ≤32×32：优化的行主序循环（L1）
/// - 任意大小：通用矩阵-向量乘法
///
/// 始终可用，不依赖任何外部库。
#[derive(Debug, Clone, Default)]
pub struct CpuBackend;

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

    /// 手写 3×3 矩阵-向量乘法（完全展开，零分支）
    #[inline(always)]
    fn mat3_mul_vec3(w: &[f32], x: &[f32], b: Option<&[f32]>) -> Vec<f32> {
        let b0 = b.map_or(0.0, |b| b[0]);
        let b1 = b.map_or(0.0, |b| b[1]);
        let b2 = b.map_or(0.0, |b| b[2]);
        vec![
            w[0] * x[0] + w[1] * x[1] + w[2] * x[2] + b0,
            w[3] * x[0] + w[4] * x[1] + w[5] * x[2] + b1,
            w[6] * x[0] + w[7] * x[1] + w[8] * x[2] + b2,
        ]
    }

    /// 通用矩阵-向量乘法（行主序）
    #[inline]
    fn general_mat_mul_vec(
        data: &[f32],
        rows: usize,
        cols: usize,
        x: &[f32],
        bias: Option<&[f32]>,
    ) -> Vec<f32> {
        let mut result = Vec::with_capacity(rows);
        for i in 0..rows {
            let row_start = i * cols;
            let mut sum = 0.0f32;
            // 展开 4 路以利于编译器向量化
            let chunks = cols / 4;
            let remainder = cols % 4;
            for c in 0..chunks {
                let base = row_start + c * 4;
                let j = c * 4;
                sum += data[base] * x[j]
                    + data[base + 1] * x[j + 1]
                    + data[base + 2] * x[j + 2]
                    + data[base + 3] * x[j + 3];
            }
            let base = row_start + chunks * 4;
            let j = chunks * 4;
            for r in 0..remainder {
                sum += data[base + r] * x[j + r];
            }
            if let Some(b) = bias {
                sum += b[i];
            }
            result.push(sum);
        }
        result
    }
}

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

    fn mat_mul_vec(
        &self,
        matrix: &TransformMatrix,
        vector: &[f32],
        bias: Option<&[f32]>,
    ) -> Vec<f32> {
        // L0: 3×3 特化路径
        if matrix.rows == 3 && matrix.cols == 3 && vector.len() == 3 {
            return Self::mat3_mul_vec3(&matrix.data, vector, bias);
        }
        // L1/通用路径
        Self::general_mat_mul_vec(&matrix.data, matrix.rows, matrix.cols, vector, bias)
    }

    fn batch_mat_mul_vec(
        &self,
        matrices: &[&TransformMatrix],
        vectors: &[&[f32]],
        biases: &[Option<&[f32]>],
    ) -> Vec<Vec<f32>> {
        // 使用 rayon 并行化大批量运算
        if matrices.len() > 64 {
            use rayon::prelude::*;
            matrices
                .par_iter()
                .zip(vectors.par_iter())
                .zip(biases.par_iter())
                .map(|((m, v), b)| self.mat_mul_vec(m, v, *b))
                .collect()
        } else {
            matrices
                .iter()
                .zip(vectors.iter())
                .zip(biases.iter())
                .map(|((m, v), b)| self.mat_mul_vec(m, v, *b))
                .collect()
        }
    }

    fn max_efficient_dim(&self) -> usize {
        64 // 超过此维度建议使用 candle 后端
    }
}

// ============================================================================
// L2: CandleCpuBackend — candle-core CPU 后端（feature-gated）
// ============================================================================

/// Candle CPU 计算后端
///
/// 使用 candle-core 的张量运算进行矩阵-向量乘法。
/// 适合中等维度（32~256）矩阵运算。
///
/// 需要 `candle` feature 开启。
#[cfg(feature = "candle")]
pub struct CandleCpuBackend {
    device: candle_core::Device,
}

#[cfg(feature = "candle")]
impl CandleCpuBackend {
    pub fn new() -> Self {
        Self {
            device: candle_core::Device::Cpu,
        }
    }

    /// 从 DeviceSelector 获取 CPU 设备（始终返回 CPU）
    pub fn from_selector(_selector: &crate::device::DeviceSelector) -> Self {
        Self::new()
    }
}

#[cfg(feature = "candle")]
impl ComputeBackend for CandleCpuBackend {
    fn name(&self) -> &str {
        "candle-cpu"
    }

    fn mat_mul_vec(
        &self,
        matrix: &TransformMatrix,
        vector: &[f32],
        bias: Option<&[f32]>,
    ) -> Vec<f32> {
        // 将 TransformMatrix 和 vector 转换为 candle Tensor
        let mat_tensor = match candle_core::Tensor::from_slice(
            &matrix.data,
            (matrix.rows, matrix.cols),
            &self.device,
        ) {
            Ok(t) => t,
            Err(_) => return CpuBackend.mat_mul_vec(matrix, vector, bias),
        };

        let vec_tensor = match candle_core::Tensor::from_slice(
            vector,
            (matrix.cols, 1),
            &self.device,
        ) {
            Ok(t) => t,
            Err(_) => return CpuBackend.mat_mul_vec(matrix, vector, bias),
        };

        // 矩阵乘法
        let result = match mat_tensor.matmul(&vec_tensor) {
            Ok(t) => t,
            Err(_) => return CpuBackend.mat_mul_vec(matrix, vector, bias),
        };

        // 提取结果
        let mut output: Vec<f32> = match result.flatten_all() {
            Ok(flat) => match flat.to_vec1::<f32>() {
                Ok(v) => v,
                Err(_) => return CpuBackend.mat_mul_vec(matrix, vector, bias),
            },
            Err(_) => return CpuBackend.mat_mul_vec(matrix, vector, bias),
        };

        // 加偏置
        if let Some(b) = bias {
            for (o, bi) in output.iter_mut().zip(b.iter()) {
                *o += bi;
            }
        }

        output
    }

    fn max_efficient_dim(&self) -> usize {
        256
    }
}

// ============================================================================
// L3: CandleCudaBackend — candle CUDA GPU 后端（feature-gated）
// ============================================================================

/// Candle CUDA GPU 计算后端
///
/// 使用 NVIDIA GPU 通过 candle CUDA 后端加速大规模矩阵运算。
/// 适合大维度（>256）和大批量运算。
///
/// 需要 `candle` feature 开启且系统有 CUDA GPU。
#[cfg(feature = "candle")]
pub struct CandleCudaBackend {
    device: candle_core::Device,
}

#[cfg(feature = "candle")]
impl CandleCudaBackend {
    /// 尝试创建 CUDA 后端，如果 GPU 不可用则返回 None
    ///
    /// **注意**: 推荐使用 `select_backend()` 或 `select_backend_with_selector()`
    /// 代替直接调用此方法，以获得统一的设备管理和降级日志。
    pub fn try_new() -> Option<Self> {
        let selector = crate::device::global_selector();
        let device = selector.candle_device();
        if matches!(device, candle_core::Device::Cpu) {
            None
        } else {
            Some(Self { device })
        }
    }

    /// 从已有的 candle Device 创建（由 DeviceSelector 提供）
    ///
    /// 调用方需确保 device 不是 CPU（否则应使用 CandleCpuBackend）。
    pub fn from_device(device: candle_core::Device) -> Self {
        Self { device }
    }
}

#[cfg(feature = "candle")]
impl ComputeBackend for CandleCudaBackend {
    fn name(&self) -> &str {
        "candle-cuda"
    }

    fn is_gpu(&self) -> bool {
        true
    }

    fn mat_mul_vec(
        &self,
        matrix: &TransformMatrix,
        vector: &[f32],
        bias: Option<&[f32]>,
    ) -> Vec<f32> {
        let mat_tensor = match candle_core::Tensor::from_slice(
            &matrix.data,
            (matrix.rows, matrix.cols),
            &self.device,
        ) {
            Ok(t) => t,
            Err(_) => return CpuBackend.mat_mul_vec(matrix, vector, bias),
        };

        let vec_tensor = match candle_core::Tensor::from_slice(
            vector,
            (matrix.cols, 1),
            &self.device,
        ) {
            Ok(t) => t,
            Err(_) => return CpuBackend.mat_mul_vec(matrix, vector, bias),
        };

        let result = match mat_tensor.matmul(&vec_tensor) {
            Ok(t) => t,
            Err(_) => return CpuBackend.mat_mul_vec(matrix, vector, bias),
        };

        // 从 GPU 拷贝回 CPU
        let cpu_result = match result.to_device(&candle_core::Device::Cpu) {
            Ok(t) => t,
            Err(_) => return CpuBackend.mat_mul_vec(matrix, vector, bias),
        };

        let mut output: Vec<f32> = match cpu_result.flatten_all() {
            Ok(flat) => match flat.to_vec1::<f32>() {
                Ok(v) => v,
                Err(_) => return CpuBackend.mat_mul_vec(matrix, vector, bias),
            },
            Err(_) => return CpuBackend.mat_mul_vec(matrix, vector, bias),
        };

        if let Some(b) = bias {
            for (o, bi) in output.iter_mut().zip(b.iter()) {
                *o += bi;
            }
        }

        output
    }

    fn batch_mat_mul_vec(
        &self,
        matrices: &[&TransformMatrix],
        vectors: &[&[f32]],
        biases: &[Option<&[f32]>],
    ) -> Vec<Vec<f32>> {
        // GPU 批量运算：尝试合并为一次大矩阵乘法
        // 如果矩阵维度不一致则回退到逐个运算
        if matrices.is_empty() {
            return Vec::new();
        }
        let first_rows = matrices[0].rows;
        let first_cols = matrices[0].cols;
        let all_same_dim = matrices.iter().all(|m| m.rows == first_rows && m.cols == first_cols);

        if !all_same_dim || matrices.len() < 4 {
            // 回退到逐个运算
            return matrices
                .iter()
                .zip(vectors.iter())
                .zip(biases.iter())
                .map(|((m, v), b)| self.mat_mul_vec(m, v, *b))
                .collect();
        }

        // 批量路径：合并 N 个矩阵为 (N*rows, cols) 的大矩阵
        // 合并 N 个向量为 (cols, N) 的大矩阵
        // 但这需要 block-diagonal 乘法，比较复杂
        // 简化为逐个调用（GPU 运算本身已足够快）
        matrices
            .iter()
            .zip(vectors.iter())
            .zip(biases.iter())
            .map(|((m, v), b)| self.mat_mul_vec(m, v, *b))
            .collect()
    }

    fn max_efficient_dim(&self) -> usize {
        usize::MAX // GPU 无实际限制
    }
}

// ============================================================================
// select_backend — 自动选择最优后端（Phase-61: 使用 DeviceSelector）
// ============================================================================

/// 根据可用硬件和矩阵维度自动选择最优计算后端
///
/// **Phase-61 重构**：使用统一 `DeviceSelector` 代替独立的 GPU 探测。
///
/// 优先级：CUDA GPU > candle CPU (大矩阵) > 纯 Rust CPU
///
/// 使用全局默认 `DeviceSelector`（auto 模式），等价于
/// `select_backend_with_selector(hint_dim, &DeviceSelector::auto())`。
pub fn select_backend(hint_dim: usize) -> Box<dyn ComputeBackend> {
    #[cfg(feature = "candle")]
    {
        let selector = crate::device::global_selector();
        return select_backend_with_selector(hint_dim, selector);
    }

    #[cfg(not(feature = "candle"))]
    {
        let _ = hint_dim;
        Box::new(CpuBackend::new())
    }
}

/// 使用自定义 `DeviceSelector` 选择最优计算后端
///
/// ## 参数
/// - `hint_dim`: 预期矩阵维度（用于 L1 vs L2 决策）
/// - `selector`: 统一设备选择器（控制 GPU 使用策略）
///
/// ## 选择逻辑
///
/// | 条件 | 选择 |
/// |---|---|
/// | GPU 可用 + DevicePreference ≠ Cpu | L3: CandleCudaBackend |
/// | candle 可用 + hint_dim > 32 | L2: CandleCpuBackend |
/// | 其他 | L0/L1: CpuBackend |
#[cfg(feature = "candle")]
pub fn select_backend_with_selector(
    hint_dim: usize,
    selector: &crate::device::DeviceSelector,
) -> Box<dyn ComputeBackend> {
    use crate::device::{DevicePreference, WorkloadHint};

    // 检查用户是否请求了 GPU
    let pref = &selector.config().device_preference;
    let force_cpu = matches!(pref, DevicePreference::Cpu);

    if !force_cpu {
        // 使用 WorkloadHint 让 DeviceSelector 做智能决策
        // 对于 compute_backend，以 hint_dim^2 估算数据量
        let estimated_items = std::cmp::max(hint_dim, 100);
        let workload = WorkloadHint::with_operation(estimated_items, hint_dim, "compute_backend");

        if selector.should_use_gpu(&workload) {
            // L3: 尝试使用 DeviceSelector 获取 GPU device
            let device = selector.candle_device();
            if !matches!(device, candle_core::Device::Cpu) {
                tracing::debug!(
                    backend = "candle-cuda",
                    hint_dim = hint_dim,
                    device = ?device,
                    "ComputeBackend: selected GPU backend via DeviceSelector"
                );
                return Box::new(CandleCudaBackend { device });
            }
        }
    }

    // L2: candle CPU（大矩阵时使用）
    if hint_dim > 32 {
        tracing::debug!(
            backend = "candle-cpu",
            hint_dim = hint_dim,
            "ComputeBackend: selected candle CPU backend"
        );
        return Box::new(CandleCpuBackend::new());
    }

    // L0/L1: 纯 Rust CPU
    Box::new(CpuBackend::new())
}

/// 创建默认的 CPU 后端（始终可用）
pub fn default_backend() -> Box<dyn ComputeBackend> {
    Box::new(CpuBackend::new())
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_cpu_backend_identity_3x3() {
        let backend = CpuBackend::new();
        let m = TransformMatrix::identity(3);
        let v = vec![1.0, 2.0, 3.0];
        let result = backend.mat_mul_vec(&m, &v, None);
        assert_eq!(result, vec![1.0, 2.0, 3.0]);
    }

    #[test]
    fn test_cpu_backend_3x3_with_bias() {
        let backend = CpuBackend::new();
        let m = TransformMatrix::identity(3);
        let v = vec![1.0, 2.0, 3.0];
        let bias = vec![0.1, 0.2, 0.3];
        let result = backend.mat_mul_vec(&m, &v, Some(&bias));
        let expected = vec![1.1, 2.2, 3.3];
        for (a, b) in result.iter().zip(expected.iter()) {
            assert!((a - b).abs() < 1e-6, "Expected {b}, got {a}");
        }
    }

    #[test]
    fn test_cpu_backend_scale_matrix() {
        let backend = CpuBackend::new();
        let m = TransformMatrix::scale(3, 2.0);
        let v = vec![1.0, 2.0, 3.0];
        let result = backend.mat_mul_vec(&m, &v, None);
        assert_eq!(result, vec![2.0, 4.0, 6.0]);
    }

    #[test]
    fn test_cpu_backend_4x4() {
        let backend = CpuBackend::new();
        let m = TransformMatrix::identity(4);
        let v = vec![1.0, 2.0, 3.0, 4.0];
        let result = backend.mat_mul_vec(&m, &v, None);
        assert_eq!(result, vec![1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn test_cpu_backend_non_square() {
        let backend = CpuBackend::new();
        // 2×3 矩阵
        let m = TransformMatrix {
            rows: 2,
            cols: 3,
            data: vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0],
        };
        let v = vec![1.0, 2.0, 3.0];
        let result = backend.mat_mul_vec(&m, &v, None);
        assert_eq!(result, vec![1.0, 2.0]);
    }

    #[test]
    fn test_cpu_backend_batch() {
        let backend = CpuBackend::new();
        let m1 = TransformMatrix::identity(3);
        let m2 = TransformMatrix::scale(3, 2.0);
        let v1 = vec![1.0, 2.0, 3.0];
        let v2 = vec![4.0, 5.0, 6.0];

        let results = backend.batch_mat_mul_vec(
            &[&m1, &m2],
            &[v1.as_slice(), v2.as_slice()],
            &[None, None],
        );

        assert_eq!(results[0], vec![1.0, 2.0, 3.0]);
        assert_eq!(results[1], vec![8.0, 10.0, 12.0]);
    }

    #[test]
    fn test_activate_batch() {
        let backend = CpuBackend::new();

        let mut values = vec![-1.0, 0.0, 0.5, 1.0];
        backend.activate_batch(&mut values, ActivationFn::ReLU);
        assert_eq!(values, vec![0.0, 0.0, 0.5, 1.0]);

        let mut values2 = vec![-1.0, 0.0, 1.0];
        backend.activate_batch(&mut values2, ActivationFn::Linear);
        assert_eq!(values2, vec![-1.0, 0.0, 1.0]);
    }

    #[test]
    fn test_apply_decay() {
        let backend = CpuBackend::new();
        let mut energies = vec![1.0, 2.0, 3.0];
        let signal_strengths = vec![0.5, 0.8, 1.0];
        let decay = 0.7;
        backend.apply_decay(&mut energies, &signal_strengths, decay);

        let expected = vec![1.0 * 0.5 * 0.7, 2.0 * 0.8 * 0.7, 3.0 * 1.0 * 0.7];
        for (a, b) in energies.iter().zip(expected.iter()) {
            assert!((a - b).abs() < 1e-10, "Expected {b}, got {a}");
        }
    }

    #[test]
    fn test_select_default_backend() {
        let backend = default_backend();
        assert_eq!(backend.name(), "cpu-rust");
        assert!(!backend.is_gpu());
    }

    #[test]
    fn test_backend_info() {
        let backend = CpuBackend::new();
        let info = backend.info();
        assert_eq!(info.name, "cpu-rust");
        assert!(!info.is_gpu);
        assert_eq!(info.max_efficient_dim, 64);
    }

    #[test]
    fn test_large_matrix_cpu() {
        let backend = CpuBackend::new();
        // 16×16 矩阵
        let dim = 16;
        let m = TransformMatrix::identity(dim);
        let v: Vec<f32> = (0..dim).map(|i| i as f32).collect();
        let result = backend.mat_mul_vec(&m, &v, None);
        assert_eq!(result, v);
    }

    #[test]
    fn test_consistency_3x3_vs_general() {
        // 验证 3×3 特化路径与通用路径结果一致
        let m = TransformMatrix::from_rows(&[
            vec![1.0, 2.0, 3.0],
            vec![4.0, 5.0, 6.0],
            vec![7.0, 8.0, 9.0],
        ]);
        let v = vec![0.5, 1.5, 2.5];
        let bias = vec![0.1, 0.2, 0.3];

        let fast_result = CpuBackend::mat3_mul_vec3(&m.data, &v, Some(&bias));
        let general_result =
            CpuBackend::general_mat_mul_vec(&m.data, m.rows, m.cols, &v, Some(&bias));

        for (a, b) in fast_result.iter().zip(general_result.iter()) {
            assert!(
                (a - b).abs() < 1e-6,
                "3×3 fast and general paths differ: {a} vs {b}"
            );
        }
    }
}
