pub mod avx2;
pub mod avx512;
pub mod neon;

use crate::quantize::scalar;

/// Runtime-dispatched 4-bit nibble dot product.
/// Picks the fastest available kernel: AVX-512 → AVX2+FMA → NEON → scalar.
#[inline]
pub fn score_4bit(query: &[f32], db_packed: &[u8]) -> f32 {
    #[cfg(target_arch = "x86_64")]
    if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") {
        return unsafe { avx512::score_4bit(query, db_packed) };
    }

    #[cfg(target_arch = "x86_64")]
    if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
        return unsafe { avx2::score_4bit(query, db_packed) };
    }

    #[cfg(target_arch = "aarch64")]
    if std::arch::is_aarch64_feature_detected!("neon") {
        return unsafe { neon::score_4bit(query, db_packed) };
    }

    scalar::score_4bit(query, db_packed)
}

/// Mixed-precision dot product: first `n4` dims at 4-bit, rest at 2-bit.
/// `n4` must be a multiple of 4 (enforced by Config::with_mixed_ratio).
/// The packed layout is: [4-bit block | 2-bit block] — see QuantPipeline::encode.
#[inline]
pub fn score_mixed(query: &[f32], db_packed: &[u8], n4: usize) -> f32 {
    use crate::quantize::scalar::{packed_bytes_4bit, packed_bytes_2bit};
    let pb4 = packed_bytes_4bit(n4);
    debug_assert!(pb4 <= db_packed.len(), "pb4={pb4} > packed.len()={}", db_packed.len());
    debug_assert_eq!(n4, pb4 * 2, "n4={n4} must be a multiple of 4");
    let n2 = query.len() - n4;
    debug_assert_eq!(packed_bytes_2bit(n2), db_packed.len() - pb4);

    let s4 = score_4bit(&query[..n4],     &db_packed[..pb4]);
    let s2 = score_2bit(&query[n4..], &db_packed[pb4..]);
    s4 + s2
}

/// Runtime-dispatched 2-bit nibble dot product.
#[inline]
pub fn score_2bit(query: &[f32], db_packed: &[u8]) -> f32 {
    #[cfg(target_arch = "x86_64")]
    if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
        return unsafe { avx2::score_2bit(query, db_packed) };
    }

    #[cfg(target_arch = "aarch64")]
    if std::arch::is_aarch64_feature_detected!("neon") {
        return unsafe { neon::score_2bit(query, db_packed) };
    }

    scalar::score_2bit(query, db_packed)
}

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

    fn make_query(dim: usize) -> Vec<f32> {
        (0..dim).map(|i| (i % 6) as f32 - 2.5).collect()
    }

    fn make_packed(dim: usize) -> Vec<u8> {
        let v: Vec<f32> = (0..dim).map(|i| (i % 5) as f32 - 1.5).collect();
        scalar::encode_4bit(&v)
    }

    #[test]
    fn dispatch_4bit_matches_scalar() {
        for dim in [8, 16, 64, 128, 512, 1536] {
            let q = make_query(dim);
            let p = make_packed(dim);
            let s = scalar::score_4bit(&q, &p);
            let d = score_4bit(&q, &p);
            assert!(
                (s - d).abs() < 1e-3,
                "dim={dim}: scalar={s:.6}, dispatch={d:.6}"
            );
        }
    }

    #[test]
    fn dispatch_uses_simd_path() {
        let q = make_query(1536);
        let p = make_packed(1536);
        let score = score_4bit(&q, &p);
        assert!(score.is_finite(), "score={score}");
    }
}
