//! Auto-adaptive compute backend for heavy operations.
//!
//! Dispatches hash computation and power-iteration PCA between CPU (BLAS) and
//! GPU (wgpu) at runtime. Auto-detection selects GPU only for discrete adapters
//! (RTX, Radeon) — integrated GPUs (Apple Silicon, Intel HD) stay on CPU because
//! Accelerate/OpenBLAS via AMX is faster for our workloads.

use ndarray::{Array1, Array2, ArrayView2};
use rand::prelude::*;
use rand_chacha::ChaCha8Rng;
use rayon::prelude::*;

#[cfg(feature = "gpu")]
use std::sync::Arc;
#[cfg(feature = "gpu")]
use crate::gpu::{GpuContext, GpuError, hash_scatter::GpuHashScatter, pca::GpuPowerIteration};

/// Compute backend — either CPU (BLAS) or GPU (wgpu).
///
/// Use `ComputeBackend::auto()` to pick the fastest option at runtime.
/// Falls back to CPU if GPU initialization fails or the adapter is integrated.
pub enum ComputeBackend {
    /// CPU with BLAS (Accelerate on macOS, OpenBLAS on Linux).
    Cpu,

    /// wgpu GPU compute (Metal/Vulkan/DX12).
    /// Best for: discrete GPUs with large datasets.
    #[cfg(feature = "gpu")]
    Gpu(GpuState),
}

/// Pre-built GPU pipelines, reusable across operations.
#[cfg(feature = "gpu")]
pub struct GpuState {
    pub ctx: Arc<GpuContext>,
    pub hasher: GpuHashScatter,
    pub pca: GpuPowerIteration,
}

#[cfg(feature = "gpu")]
impl GpuState {
    /// Build all GPU pipelines from a context. Returns error if pipeline creation fails.
    pub fn new(ctx: GpuContext) -> Result<Self, GpuError> {
        let ctx = Arc::new(ctx);
        let hasher = GpuHashScatter::new(Arc::clone(&ctx))?;
        let pca = GpuPowerIteration::new(Arc::clone(&ctx))?;
        Ok(GpuState { ctx, hasher, pca })
    }
}

impl ComputeBackend {
    /// Auto-detect the best backend. Uses GPU only for discrete adapters.
    pub fn auto() -> Self {
        #[cfg(feature = "gpu")]
        {
            if let Ok(ctx) = GpuContext::new() {
                // Only use GPU for discrete adapters (RTX, Radeon, etc.)
                // Apple Silicon integrated GPU shares memory with CPU —
                // Accelerate BLAS via AMX is faster for all our operations.
                if ctx.device_type() == wgpu::DeviceType::DiscreteGpu {
                    if let Ok(state) = GpuState::new(ctx) {
                        return ComputeBackend::Gpu(state);
                    }
                }
            }
        }
        ComputeBackend::Cpu
    }

    /// Force CPU backend.
    pub fn cpu() -> Self {
        ComputeBackend::Cpu
    }

    /// Force GPU backend. Returns None if GPU is unavailable.
    #[cfg(feature = "gpu")]
    pub fn gpu() -> Option<Self> {
        let ctx = GpuContext::new().ok()?;
        let state = GpuState::new(ctx).ok()?;
        Some(ComputeBackend::Gpu(state))
    }

    /// Whether this backend is using the GPU.
    pub fn is_gpu(&self) -> bool {
        #[cfg(feature = "gpu")]
        if let ComputeBackend::Gpu(_) = self {
            return true;
        }
        false
    }

    /// Human-readable backend description.
    pub fn description(&self) -> String {
        match self {
            ComputeBackend::Cpu => "CPU (BLAS)".to_string(),
            #[cfg(feature = "gpu")]
            ComputeBackend::Gpu(state) => format!("GPU ({})", state.ctx.info()),
        }
    }

    /// Hash N embeddings against hyperplanes. Returns bucket IDs.
    ///
    /// `data`: flat f32, shape (n × d), row-major
    /// `hp`: flat f32, shape (num_bits × d), row-major
    ///
    /// Returns Vec<u64> of hash values (one per embedding).
    pub fn hash_embeddings(
        &self,
        data: &[f32],
        n: usize,
        d: usize,
        hp: &[f32],
        num_bits: usize,
    ) -> Vec<u64> {
        match self {
            ComputeBackend::Cpu => cpu_hash_embeddings(data, n, d, hp, num_bits),
            #[cfg(feature = "gpu")]
            ComputeBackend::Gpu(state) => {
                match state.hasher.hash_batch(data, hp, n, d, num_bits) {
                    Ok(hashes) => hashes.into_iter().map(|h| h as u64).collect(),
                    Err(_) => cpu_hash_embeddings(data, n, d, hp, num_bits),
                }
            }
        }
    }

    /// Power iteration PCA on covariance matrix. Returns (components k×d flat, eigenvalues).
    ///
    /// `xtx`: flat f32 covariance matrix X^T X, shape (d × d), row-major
    pub fn power_iteration(
        &self,
        xtx: &[f32],
        d: usize,
        k: usize,
        max_iter: usize,
        seed: u64,
    ) -> (Vec<f32>, Vec<f32>) {
        match self {
            ComputeBackend::Cpu => cpu_power_iteration(xtx, d, k, max_iter, seed),
            #[cfg(feature = "gpu")]
            ComputeBackend::Gpu(state) => {
                match state.pca.power_iteration(xtx, d, k, max_iter, seed) {
                    Ok(result) => result,
                    Err(_) => cpu_power_iteration(xtx, d, k, max_iter, seed),
                }
            }
        }
    }
}

impl Default for ComputeBackend {
    fn default() -> Self {
        Self::auto()
    }
}

// =========================================================================
// CPU implementations (extracted from DensityClassifier)
// =========================================================================

/// CPU hash: BLAS matmul (n×d) @ (d×bits) + parallel bit packing.
fn cpu_hash_embeddings(
    data: &[f32],
    n: usize,
    d: usize,
    hp: &[f32],
    num_bits: usize,
) -> Vec<u64> {
    let emb_view = ArrayView2::from_shape((n, d), data).unwrap();

    // hp is row-major (num_bits × d), transpose to (d × num_bits) for matmul
    let mut hp_transposed = Array2::<f32>::zeros((d, num_bits));
    for i in 0..num_bits {
        for j in 0..d {
            hp_transposed[[j, i]] = hp[i * d + j];
        }
    }

    // BLAS-accelerated: (n × d) @ (d × bits) = (n × bits)
    let dots = emb_view.dot(&hp_transposed);

    (0..n)
        .into_par_iter()
        .map(|i| {
            let mut hash = 0u64;
            for bit_idx in 0..num_bits {
                if dots[[i, bit_idx]] > 0.0 {
                    hash |= 1 << bit_idx;
                }
            }
            hash
        })
        .collect()
}

/// CPU power iteration PCA via ndarray BLAS.
fn cpu_power_iteration(
    xtx: &[f32],
    d: usize,
    k: usize,
    max_iter: usize,
    seed: u64,
) -> (Vec<f32>, Vec<f32>) {
    let xtx_view = ArrayView2::from_shape((d, d), xtx).unwrap();
    let mut components = Vec::with_capacity(k * d);
    let mut eigenvalues = Vec::with_capacity(k);
    let mut rng = ChaCha8Rng::seed_from_u64(seed + 999);
    let mut prev_components: Vec<Array1<f32>> = Vec::with_capacity(k);

    for _comp_idx in 0..k {
        let v_data: Vec<f32> = (0..d).map(|_| rng.gen_range(-1.0f32..1.0f32)).collect();
        let mut v = Array1::from_vec(v_data);
        let norm = v.dot(&v).sqrt();
        v /= norm;

        for _ in 0..max_iter {
            let mut xtxv = xtx_view.dot(&v);

            for prev_vec in &prev_components {
                let dot = xtxv.dot(prev_vec);
                xtxv = &xtxv - &(dot * prev_vec);
            }

            let norm = xtxv.dot(&xtxv).sqrt();
            if norm > 1e-10 {
                v = &xtxv / norm;
            }
        }

        // Store component
        components.extend_from_slice(v.as_slice().unwrap());
        prev_components.push(v.clone());

        // Eigenvalue: λ = v^T (X^T X) v
        let xtxv = xtx_view.dot(&v);
        let eigenvalue = v.dot(&xtxv);
        eigenvalues.push(eigenvalue);
    }

    (components, eigenvalues)
}

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

    #[test]
    fn test_auto_detection_returns_valid_backend() {
        let backend = ComputeBackend::auto();
        // On Apple Silicon (or any machine without discrete GPU), should be CPU
        // This test just verifies auto() doesn't panic
        let desc = backend.description();
        assert!(!desc.is_empty());
    }

    #[test]
    fn test_cpu_backend_default() {
        let backend = ComputeBackend::cpu();
        assert!(!backend.is_gpu());
        assert_eq!(backend.description(), "CPU (BLAS)");
    }

    #[test]
    fn test_cpu_hash_embeddings() {
        // 4 embeddings of dim 3, 2 hyperplanes
        let data = vec![
            1.0, 0.0, 0.0,
            0.0, 1.0, 0.0,
            -1.0, 0.0, 0.0,
            0.0, -1.0, 0.0,
        ];
        // hp[0] = [1, 0, 0] — positive x
        // hp[1] = [0, 1, 0] — positive y
        let hp = vec![
            1.0, 0.0, 0.0,
            0.0, 1.0, 0.0,
        ];

        let backend = ComputeBackend::cpu();
        let hashes = backend.hash_embeddings(&data, 4, 3, &hp, 2);

        assert_eq!(hashes.len(), 4);
        // [1,0,0]: x>0 → bit0=1, y=0 → bit1=0 → hash=1
        assert_eq!(hashes[0], 0b01);
        // [0,1,0]: x=0 → bit0=0, y>0 → bit1=1 → hash=2
        assert_eq!(hashes[1], 0b10);
        // [-1,0,0]: x<0 → bit0=0, y=0 → bit1=0 → hash=0
        assert_eq!(hashes[2], 0b00);
        // [0,-1,0]: x=0 → bit0=0, y<0 → bit1=0 → hash=0
        assert_eq!(hashes[3], 0b00);
    }

    #[test]
    fn test_cpu_power_iteration_diagonal() {
        // Diagonal covariance with known eigenstructure
        let d = 4;
        let mut xtx = vec![0.0f32; d * d];
        let eigen_vals = [100.0, 50.0, 10.0, 1.0];
        for i in 0..d {
            xtx[i * d + i] = eigen_vals[i];
        }

        let backend = ComputeBackend::cpu();
        let (components, eigenvalues) = backend.power_iteration(&xtx, d, 2, 30, 42);

        assert_eq!(eigenvalues.len(), 2);
        assert!((eigenvalues[0] - 100.0).abs() < 1.0);
        assert!((eigenvalues[1] - 50.0).abs() < 1.0);

        // Components should be close to standard basis vectors
        assert_eq!(components.len(), 2 * d);
        let comp0 = &components[0..d];
        let max_idx = comp0.iter().enumerate()
            .max_by(|(_, a), (_, b)| a.abs().partial_cmp(&b.abs()).unwrap())
            .unwrap().0;
        assert!(comp0[max_idx].abs() > 0.99);
    }

    #[cfg(feature = "gpu")]
    #[test]
    fn test_gpu_state_construction() {
        // Verify GpuState can be constructed (pipelines are reusable)
        let ctx = match crate::gpu::GpuContext::new() {
            Ok(ctx) => ctx,
            Err(_) => {
                eprintln!("No GPU available, skipping test");
                return;
            }
        };
        let state = GpuState::new(ctx).unwrap();
        assert!(!state.ctx.info().is_empty());
    }

    #[cfg(feature = "gpu")]
    #[test]
    fn test_backend_dispatch_hash_consistency() {
        // Verify GPU and CPU produce identical hash results
        use rand::prelude::*;
        use rand_chacha::ChaCha8Rng;

        let dim = 64;
        let n = 500;
        let num_bits = 8;
        let mut rng = ChaCha8Rng::seed_from_u64(77);

        let mut embeddings = vec![0.0f32; n * dim];
        for i in 0..n {
            let start = i * dim;
            let mut norm = 0.0f32;
            for j in 0..dim {
                let v: f32 = rng.gen_range(-1.0..1.0);
                embeddings[start + j] = v;
                norm += v * v;
            }
            let norm = norm.sqrt();
            for j in 0..dim {
                embeddings[start + j] /= norm;
            }
        }

        let mut hyperplanes = vec![0.0f32; num_bits * dim];
        for bit in 0..num_bits {
            let start = bit * dim;
            let mut norm = 0.0f32;
            for j in 0..dim {
                let v: f32 = rng.gen_range(-1.0..1.0);
                hyperplanes[start + j] = v;
                norm += v * v;
            }
            let norm = norm.sqrt();
            for j in 0..dim {
                hyperplanes[start + j] /= norm;
            }
        }

        let cpu_hashes = ComputeBackend::cpu()
            .hash_embeddings(&embeddings, n, dim, &hyperplanes, num_bits);

        let gpu_backend = match ComputeBackend::gpu() {
            Some(b) => b,
            None => {
                eprintln!("No GPU available, skipping test");
                return;
            }
        };
        let gpu_hashes = gpu_backend
            .hash_embeddings(&embeddings, n, dim, &hyperplanes, num_bits);

        let mismatches: usize = cpu_hashes.iter().zip(&gpu_hashes)
            .filter(|(a, b)| a != b)
            .count();
        assert_eq!(mismatches, 0, "CPU vs GPU: {mismatches}/{n} hashes differ");
    }

    #[cfg(feature = "gpu")]
    #[test]
    fn test_backend_dispatch_pca_consistency() {
        // Verify GPU and CPU produce similar PCA results
        let d = 8;
        let mut xtx = vec![0.0f32; d * d];
        let eigen_vals = [100.0, 50.0, 10.0, 5.0, 2.0, 1.0, 0.5, 0.1];
        for i in 0..d {
            xtx[i * d + i] = eigen_vals[i];
        }

        let (cpu_comp, cpu_ev) = ComputeBackend::cpu()
            .power_iteration(&xtx, d, 3, 30, 42);

        let gpu_backend = match ComputeBackend::gpu() {
            Some(b) => b,
            None => {
                eprintln!("No GPU available, skipping test");
                return;
            }
        };
        let (gpu_comp, gpu_ev) = gpu_backend
            .power_iteration(&xtx, d, 3, 30, 42);

        // Eigenvalues should match closely
        for i in 0..3 {
            assert!(
                (cpu_ev[i] - gpu_ev[i]).abs() < 1.0,
                "Eigenvalue {i}: cpu={} gpu={}",
                cpu_ev[i], gpu_ev[i]
            );
        }

        // Components should point in same direction (|dot| ≈ 1)
        for comp_idx in 0..3 {
            let c = &cpu_comp[comp_idx * d..(comp_idx + 1) * d];
            let g = &gpu_comp[comp_idx * d..(comp_idx + 1) * d];
            let dot: f32 = c.iter().zip(g).map(|(a, b)| a * b).sum();
            assert!(
                dot.abs() > 0.95,
                "Component {comp_idx} alignment: {dot}"
            );
        }
    }
}
