use super::tables::{CENTROIDS_4BIT, CENTROIDS_2BIT, quantize_4bit, quantize_2bit};

/// Encode a normalized-and-rotated f32 vector to 4-bit packed bytes.
/// Input `v` must be in approximately N(0,1) scale (i.e., already multiplied by sqrt(dim)).
/// Each output byte holds 2 nibbles: byte[i] = lo_nibble | (hi_nibble << 4).
pub fn encode_4bit(v: &[f32]) -> Vec<u8> {
    debug_assert!(v.len() % 2 == 0 || !v.is_empty());
    let packed_len = (v.len() + 1) / 2;
    let mut out = vec![0u8; packed_len];
    for (i, chunk) in v.chunks(2).enumerate() {
        let lo = quantize_4bit(chunk[0]);
        let hi = if chunk.len() > 1 { quantize_4bit(chunk[1]) } else { 0 };
        out[i] = lo | (hi << 4);
    }
    out
}

/// Encode to 2-bit packed bytes. 4 values per byte.
pub fn encode_2bit(v: &[f32]) -> Vec<u8> {
    let packed_len = (v.len() + 3) / 4;
    let mut out = vec![0u8; packed_len];
    for (i, chunk) in v.chunks(4).enumerate() {
        let mut byte = 0u8;
        for (j, &x) in chunk.iter().enumerate() {
            byte |= quantize_2bit(x) << (j * 2);
        }
        out[i] = byte;
    }
    out
}

/// Decode 4-bit packed bytes back to f32 (reconstruction centroids).
/// `dim` is the number of values (may be odd for the last byte).
pub fn decode_4bit(packed: &[u8], dim: usize) -> Vec<f32> {
    let mut out = Vec::with_capacity(dim);
    for &byte in packed {
        let lo = (byte & 0x0F) as usize;
        let hi = ((byte >> 4) & 0x0F) as usize;
        out.push(CENTROIDS_4BIT[lo]);
        if out.len() < dim {
            out.push(CENTROIDS_4BIT[hi]);
        }
    }
    out.truncate(dim);
    out
}

/// Decode 2-bit packed bytes back to f32.
pub fn decode_2bit(packed: &[u8], dim: usize) -> Vec<f32> {
    let mut out = Vec::with_capacity(dim);
    for &byte in packed {
        for j in 0..4 {
            if out.len() >= dim { break; }
            let idx = ((byte >> (j * 2)) & 0x03) as usize;
            out.push(CENTROIDS_2BIT[idx]);
        }
    }
    out
}

/// Asymmetric 4-bit dot product.
///
/// Query stays in f32 (already RHDH-transformed and scaled).
/// DB vector is packed nibbles. Each nibble maps to its centroid.
///
/// score ≈ dot(query, original_db_vector)
#[inline]
pub fn score_4bit(query: &[f32], db_packed: &[u8]) -> f32 {
    debug_assert_eq!(query.len(), db_packed.len() * 2);
    let mut acc = 0.0f32;
    for (i, &byte) in db_packed.iter().enumerate() {
        let lo = (byte & 0x0F) as usize;
        let hi = ((byte >> 4) & 0x0F) as usize;
        acc += query[2 * i]     * CENTROIDS_4BIT[lo];
        acc += query[2 * i + 1] * CENTROIDS_4BIT[hi];
    }
    acc
}

/// Asymmetric 2-bit dot product.
#[inline]
pub fn score_2bit(query: &[f32], db_packed: &[u8]) -> f32 {
    debug_assert_eq!(query.len(), db_packed.len() * 4);
    let mut acc = 0.0f32;
    for (i, &byte) in db_packed.iter().enumerate() {
        for j in 0..4 {
            let idx = ((byte >> (j * 2)) & 0x03) as usize;
            acc += query[4 * i + j] * CENTROIDS_2BIT[idx];
        }
    }
    acc
}

/// Bytes required to store a packed 4-bit vector of `dim` elements.
#[inline]
pub fn packed_bytes_4bit(dim: usize) -> usize {
    (dim + 1) / 2
}

/// Bytes required to store a packed 2-bit vector of `dim` elements.
#[inline]
pub fn packed_bytes_2bit(dim: usize) -> usize {
    (dim + 3) / 4
}

#[cfg(test)]
mod tests {
    use super::*;
    #[allow(unused_imports)]
    use approx::assert_relative_eq;

    fn make_test_vec(dim: usize) -> Vec<f32> {
        // Simulate normalized values in [-3, 3] range (N(0,1) scale)
        (0..dim).map(|i| ((i as f32 % 16.0) - 7.5) * 0.4).collect()
    }

    #[test]
    fn encode_decode_4bit_roundtrip() {
        let v = make_test_vec(8);
        let packed   = encode_4bit(&v);
        let decoded  = decode_4bit(&packed, v.len());
        // Lloyd-Max outer bin [-∞, -2.40] has centroid -2.73; extreme values
        // like -3.0 incur up to ~0.27 error by design (conditional mean, not bin edge).
        for (&orig, &recon) in v.iter().zip(decoded.iter()) {
            assert!((orig - recon).abs() < 0.28, "orig={orig}, recon={recon}");
        }
    }

    #[test]
    fn encode_decode_2bit_roundtrip() {
        // Use in-range values — 2-bit Lloyd-Max outer bins are wide by design
        // (the tail is rare under N(0,1)), so extreme values have large reconstruction
        // error intentionally. Test the typical operating range instead.
        let v = vec![-1.2f32, -0.8, -0.5, -0.2, 0.2, 0.5, 0.8, 1.2];
        let packed   = encode_2bit(&v);
        let decoded  = decode_2bit(&packed, v.len());
        for (&orig, &recon) in v.iter().zip(decoded.iter()) {
            assert!((orig - recon).abs() < 0.55, "orig={orig}, recon={recon}");
        }
    }

    #[test]
    fn packed_size_4bit() {
        assert_eq!(packed_bytes_4bit(8),  4);
        assert_eq!(packed_bytes_4bit(9),  5);
        assert_eq!(packed_bytes_4bit(16), 8);
    }

    #[test]
    fn packed_size_2bit() {
        assert_eq!(packed_bytes_2bit(8),  2);
        assert_eq!(packed_bytes_2bit(9),  3);
        assert_eq!(packed_bytes_2bit(16), 4);
    }

    #[test]
    fn score_4bit_approximates_dot_product() {
        // Use values already in N(0,1) scale so centroids are close to inputs.
        let dim = 64usize;
        let v: Vec<f32> = (0..dim)
            .map(|i| if i % 2 == 0 { 0.5f32 } else { -0.5 })
            .collect();
        let packed    = encode_4bit(&v);
        let score     = score_4bit(&v, &packed);
        let exact_dot: f32 = v.iter().map(|x| x * x).sum(); // = 16.0
        let rel_err = (score - exact_dot).abs() / exact_dot;
        assert!(rel_err < 0.3, "score={score}, exact={exact_dot}, rel_err={rel_err}");
    }

    #[test]
    fn score_higher_for_similar_vectors() {
        let dim = 32usize;
        let query: Vec<f32> = vec![1.0; dim];
        let similar: Vec<f32> = vec![0.9; dim];
        let different: Vec<f32> = vec![-0.9; dim];

        let packed_similar   = encode_4bit(&similar);
        let packed_different = encode_4bit(&different);

        let s_sim = score_4bit(&query, &packed_similar);
        let s_dif = score_4bit(&query, &packed_different);
        assert!(s_sim > s_dif, "similar={s_sim}, different={s_dif}");
    }
}
