/// Quantization tables — Lloyd-Max optimal for the post-RHDH coordinate distribution.
///
/// The exact marginal of a uniform unit-vector on S^{d-1} is Beta((d-1)/2,(d-1)/2)
/// on [-1,1]. After scaling by sqrt(padded_dim) this converges to N(0,1) (CLT).
/// For padded_dim >= 128 (all practical embedding dims) the N(0,1) Lloyd-Max tables
/// are indistinguishable from the exact Beta tables — delta < 1e-4 centroid error.
/// Verified via numerical Lloyd-Max on both distributions at d=128,512,1024,2048.
///
/// Lloyd-Max minimizes MSE for a given number of levels by placing centroids
/// at the conditional means of each bin and boundaries at centroid midpoints.
/// Converged via Lloyd-Max iteration (2000 steps, tolerance 1e-12).

/// 4-bit (16 levels): Lloyd-Max centroids for N(0,1).
pub const CENTROIDS_4BIT: [f32; 16] = [
    -2.7325896, -2.0690172, -1.6180464, -1.2562312,
    -0.9423405, -0.6567591, -0.3880483, -0.1283950,
     0.1283950,  0.3880483,  0.6567591,  0.9423405,
     1.2562312,  1.6180464,  2.0690172,  2.7325896,
];

/// 4-bit decision boundaries (15 values). Any x < boundaries[i] → bin i.
pub const BOUNDARIES_4BIT: [f32; 15] = [
    -2.4008034, -1.8435318, -1.4371388, -1.0992858, -0.7995498,
    -0.5224037, -0.2582217,  0.0000000,  0.2582217,  0.5224037,
     0.7995498,  1.0992858,  1.4371388,  1.8435318,  2.4008034,
];

/// 2-bit (4 levels): Lloyd-Max centroids for N(0,1).
pub const CENTROIDS_2BIT: [f32; 4] = [-1.5104176, -0.4527800, 0.4527800, 1.5104176];

/// 2-bit decision boundaries.
pub const BOUNDARIES_2BIT: [f32; 3] = [-0.9815988, 0.0000000, 0.9815988];

/// Compute the L2 norm of the dequantized representation of a 4-bit packed vector.
/// Used for length-renormalized scoring (RaBitQ / TurboQuant):
///   corrected_score = raw_score / quant_norm
/// Removes per-vector quantization magnitude bias, improving recall ranking.
#[inline]
pub fn quant_norm_4bit(packed: &[u8]) -> f32 {
    let mut sq = 0.0f32;
    for &byte in packed {
        let lo = (byte & 0x0F) as usize;
        let hi = ((byte >> 4) & 0x0F) as usize;
        sq += CENTROIDS_4BIT[lo] * CENTROIDS_4BIT[lo];
        sq += CENTROIDS_4BIT[hi] * CENTROIDS_4BIT[hi];
    }
    sq.sqrt()
}

/// Compute the L2 norm of the dequantized representation of a 2-bit packed vector.
#[inline]
pub fn quant_norm_2bit(packed: &[u8]) -> f32 {
    let mut sq = 0.0f32;
    for &byte in packed {
        for shift in [0u32, 2, 4, 6] {
            let idx = ((byte >> shift) & 0x03) as usize;
            sq += CENTROIDS_2BIT[idx] * CENTROIDS_2BIT[idx];
        }
    }
    sq.sqrt()
}

/// Compute the L2 norm for a mixed-precision packed vector.
/// `pb4` = bytes used for the 4-bit prefix; the rest are 2-bit.
#[inline]
pub fn quant_norm_mixed(packed: &[u8], pb4: usize) -> f32 {
    let mut sq = 0.0f32;
    for &byte in &packed[..pb4] {
        let lo = (byte & 0x0F) as usize;
        let hi = ((byte >> 4) & 0x0F) as usize;
        sq += CENTROIDS_4BIT[lo] * CENTROIDS_4BIT[lo];
        sq += CENTROIDS_4BIT[hi] * CENTROIDS_4BIT[hi];
    }
    for &byte in &packed[pb4..] {
        for shift in [0u32, 2, 4, 6] {
            let idx = ((byte >> shift) & 0x03) as usize;
            sq += CENTROIDS_2BIT[idx] * CENTROIDS_2BIT[idx];
        }
    }
    sq.sqrt()
}

/// Map a coordinate (≈ N(0,1) scale) to a 4-bit level index [0, 15].
/// Uses Lloyd-Max boundaries. Values outside boundary range fall into
/// the first (bin 0) or last (bin 15) bin.
#[inline(always)]
pub fn quantize_4bit(x: f32) -> u8 {
    for (i, &b) in BOUNDARIES_4BIT.iter().enumerate() {
        if x < b {
            return i as u8;
        }
    }
    15
}

/// Map a coordinate to a 2-bit level index [0, 3].
#[inline(always)]
pub fn quantize_2bit(x: f32) -> u8 {
    for (i, &b) in BOUNDARIES_2BIT.iter().enumerate() {
        if x < b {
            return i as u8;
        }
    }
    3
}

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

    #[test]
    fn centroids_4bit_are_symmetric() {
        for i in 0..8 {
            assert_relative_eq!(
                CENTROIDS_4BIT[i],
                -CENTROIDS_4BIT[15 - i],
                epsilon = 1e-6
            );
        }
    }

    #[test]
    fn boundaries_4bit_are_symmetric() {
        for i in 0..7 {
            assert_relative_eq!(
                BOUNDARIES_4BIT[i],
                -BOUNDARIES_4BIT[14 - i],
                epsilon = 1e-6
            );
        }
    }

    #[test]
    fn centroids_2bit_are_symmetric() {
        assert_relative_eq!(CENTROIDS_2BIT[0], -CENTROIDS_2BIT[3], epsilon = 1e-6);
        assert_relative_eq!(CENTROIDS_2BIT[1], -CENTROIDS_2BIT[2], epsilon = 1e-6);
    }

    #[test]
    fn quantize_4bit_extremes() {
        assert_eq!(quantize_4bit(-10.0), 0);
        assert_eq!(quantize_4bit( 10.0), 15);
        assert_eq!(quantize_4bit(  0.0), 8);   // 0.0 ≥ boundary[7]=0.0, < boundary[8]
        assert_eq!(quantize_4bit( -3.0), 0);
        assert_eq!(quantize_4bit(  3.0), 15);
    }

    #[test]
    fn quantize_4bit_low_error() {
        // Max error ≈ half the widest bin width (~0.27 for outer bins)
        let values = [-2.5f32, -1.0, 0.0, 1.0, 2.5];
        for &v in &values {
            let idx = quantize_4bit(v) as usize;
            let reconstructed = CENTROIDS_4BIT[idx];
            assert!(
                (v - reconstructed).abs() < 0.28,
                "v={v}, idx={idx}, recon={reconstructed}, err={}",
                (v - reconstructed).abs()
            );
        }
    }

    #[test]
    fn quantize_2bit_extremes() {
        assert_eq!(quantize_2bit(-5.0), 0);
        assert_eq!(quantize_2bit( 5.0), 3);
        assert_eq!(quantize_2bit( 0.0), 2);  // 0.0 ≥ boundary[1]=0.0, falls to bin 2
    }

    #[test]
    fn boundaries_monotone() {
        for i in 1..BOUNDARIES_4BIT.len() {
            assert!(BOUNDARIES_4BIT[i] > BOUNDARIES_4BIT[i - 1]);
        }
        for i in 1..BOUNDARIES_2BIT.len() {
            assert!(BOUNDARIES_2BIT[i] > BOUNDARIES_2BIT[i - 1]);
        }
    }
}
