/// AVX2 + FMA nibble scoring kernel.
///
/// Computes the asymmetric dot product: Σ query[i] × centroid[nibble[i]]
/// using Lloyd-Max centroids (non-uniform spacing, N(0,1)-optimal).
///
/// Strategy (4-bit):
///   Process 4 packed bytes (8 nibbles = 8 f32 dimensions) per AVX2 iteration.
///   1. Extract lo/hi nibbles from 4 bytes.
///   2. Interleave to get dimension-order: [lo0, hi0, lo1, hi1, ...].
///   3. Zero-extend 8 u8 → 8 i32.
///   4. Lookup centroids using permutevar8x32 (3-cycle latency vs 13 for gather).
///      Split table into low [0..7] and high [8..15]; blend based on nibble ≥ 8.
///   5. Load 8 query f32, FMA accumulate.
///   Tail remainder handled by scalar fallback.

#[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.
/// Returns the asymmetric dot product approximation.
///
/// # Safety
/// Caller must ensure AVX2 and FMA are available (checked via `is_x86_feature_detected!`).
/// Decode one 4-byte block (8 nibbles = 8 dims) to its 8 Lloyd-Max centroid
/// values as an `__m256`. Hoisted so the main loop can run several blocks with
/// independent accumulators, hiding FMA latency.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
#[inline]
unsafe fn block_centroids(
    pb:     *const u8,
    cen_lo: __m256,
    cen_hi: __m256,
    seven:  __m256i,
    mask4:  __m128i,
) -> __m256 {
    let p = _mm_cvtsi32_si128(i32::from_le_bytes([
        *pb, *pb.add(1), *pb.add(2), *pb.add(3),
    ]));
    let lo          = _mm_and_si128(p, mask4);
    let hi          = _mm_and_si128(_mm_srli_epi16(p, 4), mask4);
    let interleaved = _mm_unpacklo_epi8(lo, hi);
    let nibble_i32  = _mm256_cvtepu8_epi32(interleaved);
    let idx_mod8    = _mm256_and_si256(nibble_i32, seven);
    let c_lo        = _mm256_permutevar8x32_ps(cen_lo, idx_mod8);
    let c_hi        = _mm256_permutevar8x32_ps(cen_hi, idx_mod8);
    let use_hi      = _mm256_cmpgt_epi32(nibble_i32, seven);
    _mm256_blendv_ps(c_lo, c_hi, _mm256_castsi256_ps(use_hi))
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
pub unsafe fn score_4bit(query: &[f32], db_packed: &[u8]) -> f32 {
    debug_assert_eq!(query.len(), db_packed.len() * 2);

    let mask4  = _mm_set1_epi8(0x0F_u8 as i8);
    // Split 16 Lloyd-Max centroids into lo [0..7] and hi [8..15] for permutevar8x32.
    let cen_lo = _mm256_loadu_ps(CENTROIDS_4BIT.as_ptr());
    let cen_hi = _mm256_loadu_ps(CENTROIDS_4BIT.as_ptr().add(8));
    let seven  = _mm256_set1_epi32(7);

    // Each block = 4 packed bytes = 8 nibbles = 8 dims.
    let n_blocks   = db_packed.len() / 4;
    let block_tail = db_packed.len() % 4;

    // Four independent accumulators to hide the ~4-cycle FMA latency: with one
    // accumulator each FMA waits on the previous (latency-bound); with four the
    // FMA throughput pipeline stays full. Summation order is fixed, so results
    // are deterministic (and within epsilon of the scalar reference).
    let mut acc0 = _mm256_setzero_ps();
    let mut acc1 = _mm256_setzero_ps();
    let mut acc2 = _mm256_setzero_ps();
    let mut acc3 = _mm256_setzero_ps();

    let n_quads = n_blocks / 4;
    for i in 0..n_quads {
        let b = i * 4;
        let c0 = block_centroids(db_packed.as_ptr().add((b    ) * 4), cen_lo, cen_hi, seven, mask4);
        let c1 = block_centroids(db_packed.as_ptr().add((b + 1) * 4), cen_lo, cen_hi, seven, mask4);
        let c2 = block_centroids(db_packed.as_ptr().add((b + 2) * 4), cen_lo, cen_hi, seven, mask4);
        let c3 = block_centroids(db_packed.as_ptr().add((b + 3) * 4), cen_lo, cen_hi, seven, mask4);
        let q0 = _mm256_loadu_ps(query.as_ptr().add((b    ) * 8));
        let q1 = _mm256_loadu_ps(query.as_ptr().add((b + 1) * 8));
        let q2 = _mm256_loadu_ps(query.as_ptr().add((b + 2) * 8));
        let q3 = _mm256_loadu_ps(query.as_ptr().add((b + 3) * 8));
        acc0 = _mm256_fmadd_ps(q0, c0, acc0);
        acc1 = _mm256_fmadd_ps(q1, c1, acc1);
        acc2 = _mm256_fmadd_ps(q2, c2, acc2);
        acc3 = _mm256_fmadd_ps(q3, c3, acc3);
    }

    // Combine the four accumulators (fixed order), then handle the 0–3 leftover blocks.
    let mut acc = _mm256_add_ps(_mm256_add_ps(acc0, acc1), _mm256_add_ps(acc2, acc3));
    for i in (n_quads * 4)..n_blocks {
        let c = block_centroids(db_packed.as_ptr().add(i * 4), cen_lo, cen_hi, seven, mask4);
        let q = _mm256_loadu_ps(query.as_ptr().add(i * 8));
        acc = _mm256_fmadd_ps(q, c, acc);
    }

    let mut result = hsum256(acc);

    // Scalar tail for the last 0–3 packed bytes (unchanged).
    let tail_byte_start = n_blocks * 4;
    for j in 0..block_tail {
        let byte_idx = tail_byte_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
}

/// Score a transformed query against a 2-bit packed DB vector using AVX2+FMA.
///
/// Strategy: process 2 packed bytes (8 2-bit indices = 8 dims) per iteration.
/// For each pair of bytes we extract 8 indices [0..3] via scalar bit ops,
/// build an AVX2 centroid vector via blendv (3 compares), then FMA
/// with 8 query floats. Scalar tail handles the last 0–1 bytes.
///
/// Centroid lookup uses 3 `blendv_ps` instead of a gather because the table
/// has only 4 entries — far cheaper than a 16-entry split-table.
///
/// # Safety
/// Caller must ensure AVX2 and FMA are available.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
pub unsafe fn score_2bit(query: &[f32], db_packed: &[u8]) -> f32 {
    use crate::quantize::tables::CENTROIDS_2BIT;
    debug_assert_eq!(query.len(), db_packed.len() * 4);

    let c0 = _mm256_set1_ps(CENTROIDS_2BIT[0]);
    let c1 = _mm256_set1_ps(CENTROIDS_2BIT[1]);
    let c2 = _mm256_set1_ps(CENTROIDS_2BIT[2]);
    let c3 = _mm256_set1_ps(CENTROIDS_2BIT[3]);
    let one = _mm256_set1_epi32(1);
    let two = _mm256_set1_epi32(2);

    let mut acc = _mm256_setzero_ps();

    // 2 packed bytes = 8 2-bit indices = 8 dims = one __m256 query register.
    let n_blocks = db_packed.len() / 2;
    let tail     = db_packed.len() % 2;

    for i in 0..n_blocks {
        let b0 = db_packed[i * 2] as i32;
        let b1 = db_packed[i * 2 + 1] as i32;

        // Extract eight 2-bit indices in dimension order.
        // byte 0 → dims base+0..3, byte 1 → dims base+4..7.
        let idx = _mm256_set_epi32(
            (b1 >> 6) & 3,
            (b1 >> 4) & 3,
            (b1 >> 2) & 3,
             b1       & 3,
            (b0 >> 6) & 3,
            (b0 >> 4) & 3,
            (b0 >> 2) & 3,
             b0       & 3,
        );

        // Centroid lookup: 3 blendv for 4 centroids (cheaper than gather).
        let ge1  = _mm256_cmpgt_epi32(idx, _mm256_setzero_si256());
        let ge2  = _mm256_cmpgt_epi32(idx, one);
        let ge3  = _mm256_cmpgt_epi32(idx, two);
        let cent = _mm256_blendv_ps(c0, c1, _mm256_castsi256_ps(ge1));
        let cent = _mm256_blendv_ps(cent, c2, _mm256_castsi256_ps(ge2));
        let cent = _mm256_blendv_ps(cent, c3, _mm256_castsi256_ps(ge3));

        let q   = _mm256_loadu_ps(query.as_ptr().add(i * 8));
        acc     = _mm256_fmadd_ps(q, cent, acc);
    }

    let mut result = hsum256(acc);

    // Scalar tail for the last odd byte (0 or 4 dims).
    if tail > 0 {
        let byte_idx = n_blocks * 2;
        let dim_base = byte_idx * 4;
        let byte     = db_packed[byte_idx];
        for j in 0..4usize {
            if dim_base + j < query.len() {
                let idx = ((byte >> (j * 2)) & 0x03) as usize;
                result += query[dim_base + j] * CENTROIDS_2BIT[idx];
            }
        }
    }

    result
}

/// Horizontal sum of an 8-lane f32 AVX2 register.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
#[inline]
unsafe fn hsum256(v: __m256) -> f32 {
    // Sum two 128-bit halves.
    let lo128 = _mm256_castps256_ps128(v);
    let hi128 = _mm256_extractf128_ps(v, 1);
    let sum4  = _mm_add_ps(lo128, hi128);

    // Sum four f32 lanes in sum4.
    let shuf  = _mm_movehdup_ps(sum4);     // [1,1,3,3]
    let sums  = _mm_add_ps(sum4, shuf);    // [0+1, _, 2+3, _]
    let hi64  = _mm_movehl_ps(sums, sums); // [2+3, _,  _,  _]
    let total = _mm_add_ss(sums, hi64);
    _mm_cvtss_f32(total)
}

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

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

    fn make_packed_4bit(dim: usize, seed: u32) -> Vec<u8> {
        let v: Vec<f32> = (0..dim).map(|i| ((i as u32 ^ seed) % 5) as f32 - 2.0).collect();
        scalar::encode_4bit(&v)
    }

    #[test]
    fn avx2_matches_scalar_4bit() {
        if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
            return; // skip if not available
        }
        for dim in [8, 16, 32, 64, 128, 1536] {
            let query  = make_query(dim, 0xAB);
            let packed = make_packed_4bit(dim, 0xCD);

            let scalar_score = scalar::score_4bit(&query, &packed);
            let avx2_score   = unsafe { score_4bit(&query, &packed) };

            let diff = (scalar_score - avx2_score).abs();
            assert!(
                diff < 1e-3,
                "dim={dim}: scalar={scalar_score:.6}, avx2={avx2_score:.6}, diff={diff:.2e}"
            );
        }
    }

    #[test]
    fn avx2_handles_non_multiple_of_8() {
        if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
            return;
        }
        // dim=10 → 5 packed bytes, not a multiple of 4
        let dim   = 10usize;
        let query = make_query(dim, 1);
        let packed = make_packed_4bit(dim, 2);

        let s = unsafe { score_4bit(&query, &packed) };
        let r = scalar::score_4bit(&query, &packed);
        assert!((s - r).abs() < 1e-3, "s={s}, r={r}");
    }

    #[test]
    fn avx2_matches_scalar_2bit() {
        if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
            return;
        }
        for dim in [8, 16, 32, 64, 128, 1536] {
            let query: Vec<f32> = (0..dim).map(|i| (i % 7) as f32 - 3.0).collect();
            let v: Vec<f32>     = (0..dim).map(|i| (i % 5) as f32 - 2.0).collect();
            let packed = scalar::encode_2bit(&v);

            let s = scalar::score_2bit(&query, &packed);
            let a = unsafe { score_2bit(&query, &packed) };
            assert!(
                (s - a).abs() < 1e-3,
                "dim={dim}: scalar={s:.6}, avx2={a:.6}, diff={:.2e}", (s-a).abs()
            );
        }
    }

    #[test]
    fn avx2_2bit_handles_odd_byte() {
        if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
            return;
        }
        // dim=12 → 3 packed bytes (odd number: 2+1)
        let dim = 12usize;
        let query: Vec<f32> = (0..dim).map(|i| i as f32 * 0.3 - 1.5).collect();
        let v: Vec<f32>     = (0..dim).map(|i| (i % 3) as f32 - 1.0).collect();
        let packed = scalar::encode_2bit(&v);

        let s = scalar::score_2bit(&query, &packed);
        let a = unsafe { score_2bit(&query, &packed) };
        assert!((s - a).abs() < 1e-3, "s={s:.6}, a={a:.6}");
    }

    #[test]
    fn avx2_score_higher_for_similar() {
        if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
            return;
        }
        let dim   = 64;
        let query = make_query(dim, 7);

        // "similar": encode a copy of query itself → high dot product
        let similar   = scalar::encode_4bit(&query);
        // "different": encode negated query → low (negative) dot product
        let neg_query: Vec<f32> = query.iter().map(|&x| -x).collect();
        let different = scalar::encode_4bit(&neg_query);

        let s_sim = unsafe { score_4bit(&query, &similar) };
        let s_dif = unsafe { score_4bit(&query, &different) };
        assert!(s_sim > s_dif, "sim={s_sim}, dif={s_dif}");
    }
}
