/// NEON nibble scoring kernel for aarch64.
///
/// Same algorithm as AVX2 but with ARM NEON intrinsics.
/// Strategy: process 8 packed bytes (16 nibbles = 16 dims) per iteration
/// using 128-bit NEON registers (4 × f32).

#[cfg(target_arch = "aarch64")]
use crate::quantize::tables::CENTROIDS_4BIT;

#[cfg(target_arch = "aarch64")]
const A4: f32 = CENTROIDS_4BIT[0];
#[cfg(target_arch = "aarch64")]
const B4: f32 = CENTROIDS_4BIT[1] - CENTROIDS_4BIT[0];

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

/// Score a transformed query against a 4-bit packed DB vector using NEON.
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn score_4bit(query: &[f32], db_packed: &[u8]) -> f32 {
    debug_assert_eq!(query.len(), db_packed.len() * 2);

    let a_v = vdupq_n_f32(A4);
    let b_v = vdupq_n_f32(B4);

    let mut acc0 = vdupq_n_f32(0.0f32);
    let mut acc1 = vdupq_n_f32(0.0f32);
    let mut acc2 = vdupq_n_f32(0.0f32);
    let mut acc3 = vdupq_n_f32(0.0f32);

    // Process 8 packed bytes (16 nibbles = 16 dims) per loop:
    //   → 4 × vfmaq_f32 (one per 4 dims, each NEON register = 4 × f32)
    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.
        let packed = vld1_u8(pb);

        // Extract lo nibbles: packed & 0x0F
        let lo_mask = vdup_n_u8(0x0F);
        let lo8     = vand_u8(packed, lo_mask);

        // Extract hi nibbles: (packed >> 4) & 0x0F
        let hi8     = vand_u8(vshr_n_u8(packed, 4), lo_mask);

        // Interleave bytes: [l0,h0, l1,h1, l2,h2, l3,h3, l4,h4, l5,h5, l6,h6, l7,h7]
        let interleaved = vzip_u8(lo8, hi8);   // vzip produces two uint8x8_t
        // interleaved.0 = [l0,h0, l1,h1, l2,h2, l3,h3]
        // interleaved.1 = [l4,h4, l5,h5, l6,h6, l7,h7]

        // Process first 8 nibbles (4+4 dims):
        // Zero-extend 4 u8 → 4 u32 → 4 f32 (lo half of interleaved.0)
        let nib0_u16 = vmovl_u8(interleaved.0);     // 8 u16
        let nib0_lo  = vmovl_u16(vget_low_u16(nib0_u16));   // 4 u32
        let nib0_hi  = vmovl_u16(vget_high_u16(nib0_u16));  // 4 u32

        let c0 = vfmaq_f32(a_v, b_v, vcvtq_f32_u32(nib0_lo));
        let c1 = vfmaq_f32(a_v, b_v, vcvtq_f32_u32(nib0_hi));

        let q0 = vld1q_f32(qp);
        let q1 = vld1q_f32(qp.add(4));
        acc0 = vfmaq_f32(acc0, q0, c0);
        acc1 = vfmaq_f32(acc1, q1, c1);

        // Process second 8 nibbles (dims 8-15):
        let nib1_u16 = vmovl_u8(interleaved.1);
        let nib1_lo  = vmovl_u16(vget_low_u16(nib1_u16));
        let nib1_hi  = vmovl_u16(vget_high_u16(nib1_u16));

        let c2 = vfmaq_f32(a_v, b_v, vcvtq_f32_u32(nib1_lo));
        let c3 = vfmaq_f32(a_v, b_v, vcvtq_f32_u32(nib1_hi));

        let q2 = vld1q_f32(qp.add(8));
        let q3 = vld1q_f32(qp.add(12));
        acc2 = vfmaq_f32(acc2, q2, c2);
        acc3 = vfmaq_f32(acc3, q3, c3);
    }

    // Horizontal sum of all accumulators.
    let sum_v = vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3));
    let mut result = vaddvq_f32(sum_v);

    // Scalar tail.
    let tail_byte_start = n_blocks * 8;
    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] * (A4 + B4 * lo as f32);
        if dim_base + 1 < query.len() {
            result += query[dim_base + 1] * (A4 + B4 * hi as f32);
        }
    }

    result
}

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn score_2bit(query: &[f32], db_packed: &[u8]) -> f32 {
    crate::quantize::scalar::score_2bit(query, db_packed)
}
