/// AVX-512F + AVX-512BW nibble scoring kernel.
///
/// Computes the asymmetric dot product: Σ query[i] × centroid[nibble[i]]
/// using Lloyd-Max centroids stored in a single 512-bit register.
///
/// Strategy (4-bit):
///   All 16 centroids fit in one __m512 (16 × f32). Centroid lookup via
///   _mm512_permutexvar_ps — no split-table needed unlike AVX2.
///   Process 8 packed bytes (16 nibbles = 16 dims) per iteration.
///   2× throughput vs AVX2 (16 dims/iter vs 8 dims/iter).

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

use crate::quantize::tables::CENTROIDS_4BIT;

/// Score a transformed query (f32[]) against a 4-bit packed DB vector.
///
/// # Safety
/// Caller must ensure avx512f and avx512bw CPU features are available.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f,avx512bw")]
pub unsafe fn score_4bit(query: &[f32], db_packed: &[u8]) -> f32 {
    debug_assert_eq!(query.len(), db_packed.len() * 2);

    // All 16 Lloyd-Max centroids in a single 512-bit register — no split-table.
    let cen      = _mm512_loadu_ps(CENTROIDS_4BIT.as_ptr());
    let mut acc  = _mm512_setzero_ps();
    let mask4    = _mm_set1_epi8(0x0F_u8 as i8);

    // Process 8 packed bytes (16 nibbles = 16 dimensions) per AVX-512 iteration.
    let n_blocks   = db_packed.len() / 8;
    let block_tail = db_packed.len() % 8;

    for i in 0..n_blocks {
        let pb = db_packed.as_ptr().add(i * 8);
        let qp = query.as_ptr().add(i * 16);

        // Load 8 packed bytes into the lower 64 bits of a 128-bit register.
        let p = _mm_loadl_epi64(pb as *const _);

        // Extract lo/hi nibbles.
        let lo = _mm_and_si128(p, mask4);
        let hi = _mm_and_si128(_mm_srli_epi16(p, 4), mask4);

        // Interleave → dimension order: [l0,h0, l1,h1, ..., l7,h7] (16 bytes).
        let interleaved = _mm_unpacklo_epi8(lo, hi);

        // Zero-extend 16 × u8 nibble indices → 16 × i32.
        let idx = _mm512_cvtepu8_epi32(interleaved);

        // Direct centroid lookup: all 16 entries fit in one __m512.
        let centroids = _mm512_permutexvar_ps(idx, cen);

        // Load 16 query f32 and FMA accumulate.
        let q = _mm512_loadu_ps(qp);
        acc = _mm512_fmadd_ps(q, centroids, acc);
    }

    // Horizontal reduction across all 16 lanes.
    let mut result = _mm512_reduce_add_ps(acc);

    // Scalar tail for remaining 0–7 packed bytes.
    let tail_start = n_blocks * 8;
    for j in 0..block_tail {
        let byte_idx = tail_start + j;
        let dim_base = byte_idx * 2;
        let lo = (db_packed[byte_idx] & 0x0F) as usize;
        let hi = ((db_packed[byte_idx] >> 4) & 0x0F) as usize;
        result += query[dim_base] * CENTROIDS_4BIT[lo];
        if dim_base + 1 < query.len() {
            result += query[dim_base + 1] * CENTROIDS_4BIT[hi];
        }
    }

    result
}

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

    fn make_query(dim: usize) -> Vec<f32> {
        (0..dim).map(|i| ((i % 7) as f32 - 3.0) * 0.5).collect()
    }

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

    #[test]
    fn avx512_matches_scalar() {
        if !is_x86_feature_detected!("avx512f") || !is_x86_feature_detected!("avx512bw") {
            return;
        }
        for dim in [16, 32, 64, 128, 384, 768, 1536] {
            let q = make_query(dim);
            let p = make_packed(dim);
            let s = scalar::score_4bit(&q, &p);
            let a = unsafe { score_4bit(&q, &p) };
            assert!(
                (s - a).abs() < 1e-3,
                "dim={dim}: scalar={s:.6}, avx512={a:.6}, diff={:.2e}", (s-a).abs()
            );
        }
    }

    #[test]
    fn avx512_handles_tail() {
        if !is_x86_feature_detected!("avx512f") || !is_x86_feature_detected!("avx512bw") {
            return;
        }
        // dim=10 → 5 packed bytes, tail = 5 % 8 = 5
        let dim = 10;
        let q = make_query(dim);
        let p = make_packed(dim);
        let s = scalar::score_4bit(&q, &p);
        let a = unsafe { score_4bit(&q, &p) };
        assert!((s - a).abs() < 1e-3, "s={s}, a={a}");
    }
}
