// Copyright 2025 SochDB Authors
//
// Licensed under the Apache License, Version 2.0

//! Unified Quantized Vector Contract
//!
//! This module provides a single abstraction for all vector quantization formats:
//! - **F32**: Full precision (baseline)
//! - **F16/BF16**: Half precision (2× compression)
//! - **I8**: 8-bit integer quantization (4× compression)
//! - **PQ**: Product quantization (32× compression)
//! - **BPS**: Block Projection Sketch (coarse filtering)
//!
//! # Architecture
//!
//! ```text
//! Query → BPS Scan (coarse) → PQ Score (refine) → I8 Rerank (exact) → Results
//!            ↓                     ↓                    ↓
//!         1000 cands            100 cands            10 results
//! ```
//!
//! # Fallback Ladder
//!
//! The pipeline automatically falls back when:
//! - PQ codebooks not trained → Skip PQ, use I8 directly
//! - I8 not available → Use F32 rerank
//! - BPS not built → Start with PQ/I8 scan
//!
//! # Usage
//!
//! ```rust,ignore
//! let scorer = UnifiedScorer::new(&config);
//! let results = scorer.search(&query, k, recall_target)?;
//! ```

use std::fmt;

/// Quantization level (ordered by compression ratio).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(u8)]
pub enum QuantLevel {
    /// Full precision f32 (1.0× compression)
    F32 = 0,
    /// Half precision f16 (2× compression)
    F16 = 1,
    /// Brain float bf16 (2× compression, better range)
    BF16 = 2,
    /// 8-bit integer (4× compression)
    I8 = 3,
    /// Product quantization (32× compression typical)
    PQ = 4,
    /// Block projection sketch (coarse filtering only)
    BPS = 5,
}

impl QuantLevel {
    /// Bytes per dimension for this level.
    pub const fn bytes_per_dim(self) -> f32 {
        match self {
            QuantLevel::F32 => 4.0,
            QuantLevel::F16 => 2.0,
            QuantLevel::BF16 => 2.0,
            QuantLevel::I8 => 1.0,
            QuantLevel::PQ => 0.125,   // ~1 byte per 8 dims (typical)
            QuantLevel::BPS => 0.0625, // ~1 byte per 16 dims
        }
    }

    /// Expected recall at this level (rough estimates).
    pub const fn expected_recall(self) -> f32 {
        match self {
            QuantLevel::F32 => 1.0,
            QuantLevel::F16 => 0.999,
            QuantLevel::BF16 => 0.998,
            QuantLevel::I8 => 0.995,
            QuantLevel::PQ => 0.90,
            QuantLevel::BPS => 0.70,
        }
    }

    /// Compute cost per vector (relative to F32 = 1.0).
    pub const fn relative_cost(self) -> f32 {
        match self {
            QuantLevel::F32 => 1.0,
            QuantLevel::F16 => 0.6,
            QuantLevel::BF16 => 0.6,
            QuantLevel::I8 => 0.3,
            QuantLevel::PQ => 0.1,
            QuantLevel::BPS => 0.05,
        }
    }
}

impl fmt::Display for QuantLevel {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            QuantLevel::F32 => write!(f, "F32"),
            QuantLevel::F16 => write!(f, "F16"),
            QuantLevel::BF16 => write!(f, "BF16"),
            QuantLevel::I8 => write!(f, "I8"),
            QuantLevel::PQ => write!(f, "PQ"),
            QuantLevel::BPS => write!(f, "BPS"),
        }
    }
}

/// Unified storage format for quantized vectors.
///
/// This is the canonical representation that all vector data flows through.
/// Each format has a specific layout optimized for its access pattern.
#[derive(Clone, Debug)]
pub enum UnifiedQuantizedVector {
    /// Full precision f32 vector.
    F32(Vec<f32>),

    /// Half precision f16 (stored as u16 bitpattern).
    F16(Vec<u16>),

    /// Brain float bf16 (stored as u16 bitpattern).
    BF16(Vec<u16>),

    /// 8-bit integer with scale and zero-point.
    I8 {
        data: Vec<i8>,
        scale: f32,
        zero_point: i8,
    },

    /// Product quantization codes.
    PQ {
        /// PQ codes (one byte per subspace).
        codes: Vec<u8>,
        /// Number of subspaces.
        num_subspaces: usize,
        /// Precomputed scale for reconstruction.
        scale: f32,
    },

    /// Block projection sketch (for coarse filtering).
    BPS {
        /// Sketch bytes (one per block).
        sketch: Vec<u8>,
        /// Number of blocks.
        num_blocks: usize,
    },
}

/// Errors produced by quantized-vector construction and decoding.
///
/// These exist to make unsupported operations *fail fast* with a typed error
/// instead of silently returning all-zero (origin) vectors, which would make
/// distance computations meaningless and recall effectively random.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum QuantError {
    /// Product quantization requires a trained codebook that this entry point
    /// does not have access to. Train a codebook and use the codebook-aware
    /// encode/decode path instead.
    PqRequiresCodebook,
    /// The format is a coarse sketch and cannot be reconstructed to f32.
    NotReconstructable(QuantLevel),
}

impl fmt::Display for QuantError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            QuantError::PqRequiresCodebook => write!(
                f,
                "product quantization requires a trained codebook (no codebook available on this path)"
            ),
            QuantError::NotReconstructable(level) => {
                write!(
                    f,
                    "{} is a coarse sketch and cannot be reconstructed to f32",
                    level
                )
            }
        }
    }
}

impl std::error::Error for QuantError {}

impl UnifiedQuantizedVector {
    /// Get the quantization level.
    pub fn level(&self) -> QuantLevel {
        match self {
            UnifiedQuantizedVector::F32(_) => QuantLevel::F32,
            UnifiedQuantizedVector::F16(_) => QuantLevel::F16,
            UnifiedQuantizedVector::BF16(_) => QuantLevel::BF16,
            UnifiedQuantizedVector::I8 { .. } => QuantLevel::I8,
            UnifiedQuantizedVector::PQ { .. } => QuantLevel::PQ,
            UnifiedQuantizedVector::BPS { .. } => QuantLevel::BPS,
        }
    }

    /// Get dimension (or approximate for compressed formats).
    pub fn dimension(&self) -> usize {
        match self {
            UnifiedQuantizedVector::F32(v) => v.len(),
            UnifiedQuantizedVector::F16(v) => v.len(),
            UnifiedQuantizedVector::BF16(v) => v.len(),
            UnifiedQuantizedVector::I8 { data, .. } => data.len(),
            UnifiedQuantizedVector::PQ { num_subspaces, .. } => *num_subspaces * 8, // Approximate
            UnifiedQuantizedVector::BPS { num_blocks, .. } => *num_blocks * 16,     // Approximate
        }
    }

    /// Memory usage in bytes.
    pub fn memory_bytes(&self) -> usize {
        match self {
            UnifiedQuantizedVector::F32(v) => v.len() * 4,
            UnifiedQuantizedVector::F16(v) => v.len() * 2,
            UnifiedQuantizedVector::BF16(v) => v.len() * 2,
            UnifiedQuantizedVector::I8 { data, .. } => data.len() + 5, // data + scale + zero
            UnifiedQuantizedVector::PQ { codes, .. } => codes.len() + 8, // codes + metadata
            UnifiedQuantizedVector::BPS { sketch, .. } => sketch.len() + 4,
        }
    }

    /// Convert to f32 vector (decode).
    ///
    /// Returns an error for formats that cannot be faithfully reconstructed on
    /// this path (`PQ` without a codebook, and `BPS` sketches), rather than
    /// silently returning an all-zero vector.
    pub fn to_f32(&self) -> Result<Vec<f32>, QuantError> {
        match self {
            UnifiedQuantizedVector::F32(v) => Ok(v.clone()),
            UnifiedQuantizedVector::F16(v) => Ok(v.iter().map(|&x| f16_to_f32(x)).collect()),
            UnifiedQuantizedVector::BF16(v) => Ok(v.iter().map(|&x| bf16_to_f32(x)).collect()),
            UnifiedQuantizedVector::I8 {
                data,
                scale,
                zero_point,
            } => Ok(data
                .iter()
                .map(|&x| (x as f32 - *zero_point as f32) * scale)
                .collect()),
            UnifiedQuantizedVector::PQ { .. } => Err(QuantError::PqRequiresCodebook),
            UnifiedQuantizedVector::BPS { .. } => {
                Err(QuantError::NotReconstructable(QuantLevel::BPS))
            }
        }
    }

    /// Create from f32 vector with specified quantization level.
    ///
    /// Returns [`QuantError::PqRequiresCodebook`] for `QuantLevel::PQ`: product
    /// quantization cannot be performed here because there is no trained
    /// codebook on this path. Previously this produced an all-zero PQ vector,
    /// which silently corrupted distance computations.
    pub fn from_f32(data: &[f32], level: QuantLevel) -> Result<Self, QuantError> {
        match level {
            QuantLevel::F32 => Ok(UnifiedQuantizedVector::F32(data.to_vec())),
            QuantLevel::F16 => Ok(UnifiedQuantizedVector::F16(
                data.iter().map(|&x| f32_to_f16(x)).collect(),
            )),
            QuantLevel::BF16 => Ok(UnifiedQuantizedVector::BF16(
                data.iter().map(|&x| f32_to_bf16(x)).collect(),
            )),
            QuantLevel::I8 => {
                // Simple symmetric quantization
                let max_abs = data.iter().map(|x| x.abs()).fold(0.0f32, f32::max);
                let scale = if max_abs > 0.0 { max_abs / 127.0 } else { 1.0 };
                let quantized: Vec<i8> = data
                    .iter()
                    .map(|&x| (x / scale).clamp(-127.0, 127.0) as i8)
                    .collect();
                Ok(UnifiedQuantizedVector::I8 {
                    data: quantized,
                    scale,
                    zero_point: 0,
                })
            }
            QuantLevel::PQ => Err(QuantError::PqRequiresCodebook),
            QuantLevel::BPS => {
                // BPS projection
                let num_blocks = (data.len() + 15) / 16;
                let sketch: Vec<u8> = (0..num_blocks)
                    .map(|b| {
                        let start = b * 16;
                        let end = (start + 16).min(data.len());
                        let sum: f32 = data[start..end].iter().sum();
                        ((sum * 10.0).clamp(0.0, 255.0)) as u8
                    })
                    .collect();
                Ok(UnifiedQuantizedVector::BPS { sketch, num_blocks })
            }
        }
    }

    /// Encode an f32 vector into PQ codes using a trained codebook.
    ///
    /// This is the codebook-aware counterpart to [`from_f32`](Self::from_f32) for
    /// `QuantLevel::PQ`. Unlike the codebook-less path (which fails fast with
    /// [`QuantError::PqRequiresCodebook`]), this performs real product
    /// quantization: each subspace is mapped to its nearest centroid in the
    /// trained codebook, yielding one byte per subspace (~32× compression).
    ///
    /// The vector dimension must match the codebook's `original_dim`.
    pub fn from_f32_pq(data: &[f32], codebook: &crate::product_quantization::PQCodebooks) -> Self {
        let pq = codebook.encode_slice(data);
        UnifiedQuantizedVector::PQ {
            codes: pq.codes,
            num_subspaces: codebook.n_subspaces,
            scale: 1.0,
        }
    }

    /// Decode PQ codes back to an approximate f32 vector using the trained codebook.
    ///
    /// This is the codebook-aware counterpart to [`to_f32`](Self::to_f32): given the
    /// same codebook used to encode, it reconstructs the approximate vector by
    /// concatenating the per-subspace centroids. Returns
    /// [`QuantError::NotReconstructable`] if called on a non-`PQ` value.
    pub fn to_f32_pq(
        &self,
        codebook: &crate::product_quantization::PQCodebooks,
    ) -> Result<Vec<f32>, QuantError> {
        match self {
            UnifiedQuantizedVector::PQ { codes, .. } => {
                let pq = crate::product_quantization::PQCodes::from_bytes(codes);
                Ok(codebook.decode(&pq).to_vec())
            }
            other => Err(QuantError::NotReconstructable(other.level())),
        }
    }
}

/// Stage in the multi-stage retrieval pipeline.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PipelineStage {
    /// Coarse filtering (BPS scan).
    Coarse,
    /// Refinement (PQ scoring).
    Refine,
    /// Final reranking (I8 or F32 exact).
    Rerank,
}

/// Configuration for the quantization pipeline.
#[derive(Debug, Clone)]
pub struct QuantPipelineConfig {
    /// Target recall (0.0 to 1.0).
    pub target_recall: f32,

    /// Maximum latency budget in microseconds.
    pub latency_budget_us: u64,

    /// Candidate count at each stage.
    pub stage_candidates: StageCandidates,

    /// Available quantization levels.
    pub available_levels: Vec<QuantLevel>,

    /// Whether to use adaptive stage selection.
    pub adaptive: bool,
}

/// Candidate counts at each pipeline stage.
#[derive(Debug, Clone)]
pub struct StageCandidates {
    /// Candidates after coarse stage.
    pub after_coarse: usize,
    /// Candidates after refinement stage.
    pub after_refine: usize,
    /// Final k results.
    pub final_k: usize,
}

impl Default for QuantPipelineConfig {
    fn default() -> Self {
        Self {
            target_recall: 0.95,
            latency_budget_us: 1000, // 1ms
            stage_candidates: StageCandidates {
                after_coarse: 1000,
                after_refine: 100,
                final_k: 10,
            },
            available_levels: vec![QuantLevel::F32, QuantLevel::I8],
            adaptive: true,
        }
    }
}

/// Unified scorer that handles all quantization formats.
pub struct UnifiedScorer {
    /// Pipeline configuration.
    config: QuantPipelineConfig,

    /// Best available level for reranking.
    rerank_level: QuantLevel,
}

impl UnifiedScorer {
    /// Create a new unified scorer.
    pub fn new(config: QuantPipelineConfig) -> Self {
        // Determine best rerank level from available
        let rerank_level = config
            .available_levels
            .iter()
            .filter(|l| matches!(l, QuantLevel::F32 | QuantLevel::I8))
            .min()
            .copied()
            .unwrap_or(QuantLevel::F32);

        Self {
            config,
            rerank_level,
        }
    }

    /// Get the rerank level.
    pub fn rerank_level(&self) -> QuantLevel {
        self.rerank_level
    }

    /// Compute I8 dot product between query and vector.
    #[inline]
    pub fn dot_i8(query: &[i8], vector: &[i8]) -> i32 {
        // Use SIMD if available
        #[cfg(target_arch = "x86_64")]
        {
            if is_x86_feature_detected!("avx2") {
                return unsafe { Self::dot_i8_avx2(query, vector) };
            }
        }

        // Scalar fallback
        query
            .iter()
            .zip(vector.iter())
            .map(|(&a, &b)| (a as i32) * (b as i32))
            .sum()
    }

    /// Compute I8 dot product with dequantization.
    #[inline]
    pub fn dot_i8_dequant(query: &[i8], query_scale: f32, vector: &[i8], vector_scale: f32) -> f32 {
        let int_dot = Self::dot_i8(query, vector);
        int_dot as f32 * query_scale * vector_scale
    }

    /// AVX2 implementation of I8 dot product.
    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx2")]
    unsafe fn dot_i8_avx2(a: &[i8], b: &[i8]) -> i32 {
        use std::arch::x86_64::*;

        let len = a.len().min(b.len());
        let chunks = len / 32;

        let mut acc = _mm256_setzero_si256();

        for i in 0..chunks {
            let idx = i * 32;
            let va = _mm256_loadu_si256(a.as_ptr().add(idx) as *const __m256i);
            let vb = _mm256_loadu_si256(b.as_ptr().add(idx) as *const __m256i);

            // Sign-extend and multiply-add
            let va_lo = _mm256_castsi256_si128(va);
            let va_hi = _mm256_extracti128_si256(va, 1);
            let vb_lo = _mm256_castsi256_si128(vb);
            let vb_hi = _mm256_extracti128_si256(vb, 1);

            let a_lo_16 = _mm256_cvtepi8_epi16(va_lo);
            let a_hi_16 = _mm256_cvtepi8_epi16(va_hi);
            let b_lo_16 = _mm256_cvtepi8_epi16(vb_lo);
            let b_hi_16 = _mm256_cvtepi8_epi16(vb_hi);

            let prod_lo = _mm256_madd_epi16(a_lo_16, b_lo_16);
            let prod_hi = _mm256_madd_epi16(a_hi_16, b_hi_16);

            acc = _mm256_add_epi32(acc, prod_lo);
            acc = _mm256_add_epi32(acc, prod_hi);
        }

        // Horizontal sum
        let acc_lo = _mm256_castsi256_si128(acc);
        let acc_hi = _mm256_extracti128_si256(acc, 1);
        let sum128 = _mm_add_epi32(acc_lo, acc_hi);
        let sum128 = _mm_hadd_epi32(sum128, sum128);
        let sum128 = _mm_hadd_epi32(sum128, sum128);
        let mut total = _mm_cvtsi128_si32(sum128);

        // Handle remainder
        for i in (chunks * 32)..len {
            total += (a[i] as i32) * (b[i] as i32);
        }

        total
    }

    /// Estimate recall at given candidate count for a level.
    pub fn estimate_recall(
        &self,
        level: QuantLevel,
        candidates: usize,
        total_vectors: usize,
    ) -> f32 {
        let base_recall = level.expected_recall();
        let coverage = (candidates as f32 / total_vectors as f32).min(1.0);
        base_recall * coverage.sqrt() // Rough model
    }

    /// Choose optimal pipeline stages given constraints.
    pub fn plan_pipeline(&self, total_vectors: usize) -> Vec<(PipelineStage, QuantLevel, usize)> {
        let mut stages = Vec::new();

        let has_bps = self.config.available_levels.contains(&QuantLevel::BPS);
        let has_pq = self.config.available_levels.contains(&QuantLevel::PQ);

        // Coarse stage (if BPS available and dataset is large enough)
        if has_bps && total_vectors > 10_000 {
            stages.push((
                PipelineStage::Coarse,
                QuantLevel::BPS,
                self.config.stage_candidates.after_coarse,
            ));
        }

        // Refine stage (if PQ available)
        if has_pq && total_vectors > 1_000 {
            stages.push((
                PipelineStage::Refine,
                QuantLevel::PQ,
                self.config.stage_candidates.after_refine,
            ));
        }

        // Rerank stage (always)
        stages.push((
            PipelineStage::Rerank,
            self.rerank_level,
            self.config.stage_candidates.final_k,
        ));

        stages
    }
}

// ============================================================================
// F16/BF16 conversion utilities
// ============================================================================

// f16/bf16 conversions delegate to the `half` crate (already a dependency,
// also used by vector_quantized.rs/compression.rs). The previous hand-rolled
// converters mangled NaN: `f32_to_f16` did `0x7C00 | (frac >> 13)`, turning
// any NaN whose payload lived in the low 13 bits (e.g. 0x7F800001) into +Inf,
// and `f32_to_bf16`'s plain `>> 16` truncation turned low-payload NaN into
// +Inf likewise. half's conversions preserve NaN and round correctly.

/// Convert f32 to f16 (IEEE 754 half-precision).
#[inline]
fn f32_to_f16(x: f32) -> u16 {
    half::f16::from_f32(x).to_bits()
}

/// Convert f16 to f32.
#[inline]
fn f16_to_f32(x: u16) -> f32 {
    half::f16::from_bits(x).to_f32()
}

/// Convert f32 to bf16 (Brain float).
#[inline]
fn f32_to_bf16(x: f32) -> u16 {
    half::bf16::from_f32(x).to_bits()
}

/// Convert bf16 to f32.
#[inline]
fn bf16_to_f32(x: u16) -> f32 {
    half::bf16::from_bits(x).to_f32()
}

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

    #[test]
    fn test_quant_level_ordering() {
        assert!(QuantLevel::F32 < QuantLevel::I8);
        assert!(QuantLevel::I8 < QuantLevel::PQ);
        assert!(QuantLevel::PQ < QuantLevel::BPS);
    }

    #[test]
    fn test_f16_bf16_preserve_nan_not_infinity() {
        // Regression: the old hand-rolled converters turned a NaN whose payload
        // lived in the low bits (e.g. 0x7F800001) into +Infinity. half-crate
        // delegation must keep NaN as NaN for both f16 and bf16.
        for payload in [0x7F800001u32, 0x7F801000, 0x7FC00000, 0xFF800042] {
            let x = f32::from_bits(payload);
            assert!(x.is_nan(), "test input must be NaN");
            assert!(
                f16_to_f32(f32_to_f16(x)).is_nan(),
                "f16 round-trip turned NaN 0x{:08X} into {:?}",
                payload,
                f16_to_f32(f32_to_f16(x))
            );
            assert!(
                bf16_to_f32(f32_to_bf16(x)).is_nan(),
                "bf16 round-trip turned NaN 0x{:08X} into {:?}",
                payload,
                bf16_to_f32(f32_to_bf16(x))
            );
        }
        // Sanity: finite + inf still round-trip sensibly.
        assert_eq!(f16_to_f32(f32_to_f16(1.5)), 1.5);
        assert!(f16_to_f32(f32_to_f16(f32::INFINITY)).is_infinite());
        assert!(bf16_to_f32(f32_to_bf16(f32::NEG_INFINITY)).is_infinite());
    }

    #[test]
    fn test_unified_vector_f32() {
        let data = vec![1.0, 2.0, 3.0, 4.0];
        let vec = UnifiedQuantizedVector::from_f32(&data, QuantLevel::F32).unwrap();

        assert_eq!(vec.level(), QuantLevel::F32);
        assert_eq!(vec.dimension(), 4);
        assert_eq!(vec.to_f32().unwrap(), data);
    }

    #[test]
    fn test_unified_vector_i8() {
        let data = vec![0.5, -0.3, 0.8, -0.1];
        let vec = UnifiedQuantizedVector::from_f32(&data, QuantLevel::I8).unwrap();

        assert_eq!(vec.level(), QuantLevel::I8);
        assert_eq!(vec.dimension(), 4);

        // Check reconstruction accuracy
        let reconstructed = vec.to_f32().unwrap();
        for (orig, recon) in data.iter().zip(reconstructed.iter()) {
            assert!(
                (orig - recon).abs() < 0.1,
                "I8 reconstruction error too large"
            );
        }
    }

    #[test]
    fn test_pq_fails_fast_without_codebook() {
        // PQ construction must fail fast rather than produce all-zero codes.
        let data = vec![0.5, -0.3, 0.8, -0.1, 0.2, 0.9, -0.4, 0.1];
        let err = UnifiedQuantizedVector::from_f32(&data, QuantLevel::PQ).unwrap_err();
        assert_eq!(err, QuantError::PqRequiresCodebook);
    }

    #[test]
    fn test_bps_is_not_reconstructable() {
        // A coarse sketch must not silently decode to zeros.
        let data = vec![0.5, -0.3, 0.8, -0.1];
        let sketch = UnifiedQuantizedVector::from_f32(&data, QuantLevel::BPS).unwrap();
        let err = sketch.to_f32().unwrap_err();
        assert_eq!(err, QuantError::NotReconstructable(QuantLevel::BPS));
    }

    #[test]
    fn test_pq_roundtrip_with_codebook() {
        use crate::product_quantization::PQCodebooks;
        use ndarray::Array1;

        // Train a tiny codebook on synthetic 16-dim data (subdim=8 → 2 subspaces).
        let mut training: Vec<Array1<f32>> = Vec::new();
        for i in 0..256 {
            let base = (i % 8) as f32;
            training.push(Array1::from_vec(
                (0..16)
                    .map(|d| base + 0.01 * (d as f32) + 0.001 * (i as f32))
                    .collect(),
            ));
        }
        let codebook = PQCodebooks::train(&training, 5, 8);

        // Encode a vector that exists in the training distribution.
        let original: Vec<f32> = (0..16).map(|d| 3.0 + 0.01 * (d as f32)).collect();
        let pq = UnifiedQuantizedVector::from_f32_pq(&original, &codebook);
        assert_eq!(pq.level(), QuantLevel::PQ);

        // Codebook-aware decode must reconstruct an approximation (not zeros).
        let recon = pq.to_f32_pq(&codebook).unwrap();
        assert_eq!(recon.len(), original.len());
        let nonzero = recon.iter().any(|&x| x.abs() > 1e-6);
        assert!(nonzero, "PQ reconstruction must not be all zeros");

        // Reconstruction should be reasonably close to the original for in-distribution data.
        let mse: f32 = original
            .iter()
            .zip(recon.iter())
            .map(|(o, r)| (o - r).powi(2))
            .sum::<f32>()
            / original.len() as f32;
        assert!(mse < 1.0, "PQ reconstruction MSE too large: {mse}");
    }

    #[test]
    fn test_to_f32_pq_rejects_non_pq() {
        use crate::product_quantization::PQCodebooks;
        use ndarray::Array1;

        let training: Vec<Array1<f32>> = (0..256)
            .map(|i| Array1::from_vec((0..16).map(|d| (i + d) as f32).collect()))
            .collect();
        let codebook = PQCodebooks::train(&training, 3, 8);

        // An F32 value is not PQ-encoded; codebook decode must reject it.
        let f32_vec = UnifiedQuantizedVector::from_f32(&[1.0; 16], QuantLevel::F32).unwrap();
        let err = f32_vec.to_f32_pq(&codebook).unwrap_err();
        assert_eq!(err, QuantError::NotReconstructable(QuantLevel::F32));
    }

    #[test]
    fn test_unified_vector_f16() {
        let data = vec![1.5, -2.25, 0.0, 100.0];
        let vec = UnifiedQuantizedVector::from_f32(&data, QuantLevel::F16).unwrap();

        assert_eq!(vec.level(), QuantLevel::F16);

        let reconstructed = vec.to_f32().unwrap();
        for (orig, recon) in data.iter().zip(reconstructed.iter()) {
            assert!(
                (orig - recon).abs() < 0.01 * orig.abs().max(1.0),
                "F16 reconstruction error too large"
            );
        }
    }

    #[test]
    fn test_unified_vector_bf16() {
        let data = vec![1.5, -2.25, 0.0, 100.0];
        let vec = UnifiedQuantizedVector::from_f32(&data, QuantLevel::BF16).unwrap();

        assert_eq!(vec.level(), QuantLevel::BF16);

        let reconstructed = vec.to_f32().unwrap();
        for (orig, recon) in data.iter().zip(reconstructed.iter()) {
            // BF16 has lower precision than F16
            assert!(
                (orig - recon).abs() < 0.1 * orig.abs().max(1.0),
                "BF16 reconstruction error too large"
            );
        }
    }

    #[test]
    fn test_dot_i8() {
        let a: Vec<i8> = vec![1, 2, 3, 4, 5, 6, 7, 8];
        let b: Vec<i8> = vec![8, 7, 6, 5, 4, 3, 2, 1];

        let result = UnifiedScorer::dot_i8(&a, &b);
        let expected: i32 = a
            .iter()
            .zip(b.iter())
            .map(|(&x, &y)| (x as i32) * (y as i32))
            .sum();

        assert_eq!(result, expected);
    }

    #[test]
    fn test_pipeline_planning() {
        let config = QuantPipelineConfig {
            available_levels: vec![
                QuantLevel::F32,
                QuantLevel::I8,
                QuantLevel::BPS,
                QuantLevel::PQ,
            ],
            ..Default::default()
        };

        let scorer = UnifiedScorer::new(config);

        // Small dataset: should skip coarse and refine
        let stages_small = scorer.plan_pipeline(500);
        assert_eq!(stages_small.len(), 1);
        assert_eq!(stages_small[0].0, PipelineStage::Rerank);

        // Large dataset: should use all stages
        let stages_large = scorer.plan_pipeline(100_000);
        assert!(stages_large.len() >= 2);
    }

    #[test]
    fn test_f16_roundtrip() {
        let values = [0.0, 1.0, -1.0, 0.5, 100.0, -0.001, 65504.0]; // Max F16 value

        for &v in &values {
            let f16_bits = f32_to_f16(v);
            let back = f16_to_f32(f16_bits);

            if v == 0.0 {
                assert_eq!(back, 0.0);
            } else {
                let rel_error = ((v - back) / v).abs();
                assert!(
                    rel_error < 0.001 || (v - back).abs() < 0.001,
                    "F16 roundtrip failed for {}: got {}",
                    v,
                    back
                );
            }
        }
    }

    #[test]
    fn test_bf16_roundtrip() {
        let values = [0.0, 1.0, -1.0, 100.0, 1e10, -1e-10];

        for &v in &values {
            let bf16_bits = f32_to_bf16(v);
            let back = bf16_to_f32(bf16_bits);

            if v == 0.0 {
                assert!(back.abs() < 1e-10);
            } else {
                let rel_error = ((v - back) / v).abs();
                assert!(
                    rel_error < 0.01,
                    "BF16 roundtrip failed for {}: got {} (error {})",
                    v,
                    back,
                    rel_error
                );
            }
        }
    }

    #[test]
    fn test_memory_usage() {
        let dim = 768;
        let data: Vec<f32> = (0..dim).map(|i| i as f32 / 1000.0).collect();

        let f32_vec = UnifiedQuantizedVector::from_f32(&data, QuantLevel::F32).unwrap();
        let i8_vec = UnifiedQuantizedVector::from_f32(&data, QuantLevel::I8).unwrap();
        let f16_vec = UnifiedQuantizedVector::from_f32(&data, QuantLevel::F16).unwrap();

        assert_eq!(f32_vec.memory_bytes(), dim * 4);
        assert_eq!(i8_vec.memory_bytes(), dim + 5); // data + scale + zero
        assert_eq!(f16_vec.memory_bytes(), dim * 2);

        // Verify compression ratios
        assert!(i8_vec.memory_bytes() < f32_vec.memory_bytes() / 3);
        assert!(f16_vec.memory_bytes() < f32_vec.memory_bytes());
    }
}
