//! Dequantization: convert GGML quantized weights to f32.
//!
//! Supports: F32, F16, Q8_0, Q5_0, Q4_K_M (Q4K), Q6_K (Q6K).

use crate::gguf::GgmlType;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum QuantError {
    #[error("Unsupported quantization type: {0:?}")]
    UnsupportedType(GgmlType),
    #[error("Invalid block size: expected {expected}, got {got}")]
    InvalidBlockSize { expected: usize, got: usize },
}

pub fn block_size_bytes(dtype: GgmlType) -> Option<usize> {
    match dtype {
        GgmlType::F32 => Some(4),
        GgmlType::F16 => Some(2),
        GgmlType::Q8_0 => Some(34),   // 2 (f16 scale) + 32 (i8 quants)
        GgmlType::Q5_0 => Some(22),   // 2 (f16 scale) + 4 (high bits u32) + 16 (32 nibbles)
        GgmlType::Q4K => Some(144),   // 2+2+12+128 = 144
        GgmlType::Q6K => Some(210),   // 128+64+16+2 = 210
        _ => None,
    }
}

pub fn block_elements(dtype: GgmlType) -> Option<usize> {
    match dtype {
        GgmlType::F32 => Some(1),
        GgmlType::F16 => Some(1),
        GgmlType::Q8_0 => Some(32),
        GgmlType::Q5_0 => Some(32),
        GgmlType::Q4K => Some(256),
        GgmlType::Q6K => Some(256),
        _ => None,
    }
}

pub fn dequantize(data: &[u8], dtype: GgmlType, n_elements: usize) -> Result<Vec<f32>, QuantError> {
    match dtype {
        GgmlType::F32 => dequant_f32(data, n_elements),
        GgmlType::F16 => dequant_f16(data, n_elements),
        GgmlType::Q8_0 => dequant_q8_0(data, n_elements),
        GgmlType::Q5_0 => dequant_q5_0(data, n_elements),
        GgmlType::Q4K => dequant_q4_k(data, n_elements),
        GgmlType::Q6K => dequant_q6_k(data, n_elements),
        other => Err(QuantError::UnsupportedType(other)),
    }
}

fn dequant_f32(data: &[u8], n: usize) -> Result<Vec<f32>, QuantError> {
    if data.len() < n * 4 {
        return Err(QuantError::InvalidBlockSize {
            expected: n * 4,
            got: data.len(),
        });
    }
    let mut out = Vec::with_capacity(n);
    for i in 0..n {
        let bytes = [
            data[i * 4],
            data[i * 4 + 1],
            data[i * 4 + 2],
            data[i * 4 + 3],
        ];
        out.push(f32::from_le_bytes(bytes));
    }
    Ok(out)
}

fn dequant_f16(data: &[u8], n: usize) -> Result<Vec<f32>, QuantError> {
    if data.len() < n * 2 {
        return Err(QuantError::InvalidBlockSize {
            expected: n * 2,
            got: data.len(),
        });
    }
    let mut out = Vec::with_capacity(n);
    for i in 0..n {
        let bytes = [data[i * 2], data[i * 2 + 1]];
        let h = half::f16::from_le_bytes(bytes);
        out.push(h.to_f32());
    }
    Ok(out)
}

fn dequant_q8_0(data: &[u8], n_elements: usize) -> Result<Vec<f32>, QuantError> {
    let block_size = 34;
    let elements_per_block = 32;
    let n_blocks = n_elements.div_ceil(elements_per_block);
    if data.len() < n_blocks * block_size {
        return Err(QuantError::InvalidBlockSize {
            expected: n_blocks * block_size,
            got: data.len(),
        });
    }
    let mut out = Vec::with_capacity(n_elements);
    for block_idx in 0..n_blocks {
        let block_start = block_idx * block_size;
        let scale_bytes = [data[block_start], data[block_start + 1]];
        let scale = half::f16::from_le_bytes(scale_bytes).to_f32();
        for j in 0..elements_per_block {
            if out.len() >= n_elements {
                break;
            }
            let q = data[block_start + 2 + j] as i8;
            out.push(q as f32 * scale);
        }
    }
    Ok(out)
}

/// Q5_0: 32 elements per block, 22 bytes.
///
/// Block layout:
///   - 2 bytes: f16 scale (d)
///   - 4 bytes: high-bit mask (bit 4 of each element, packed as u32 LE)
///   - 16 bytes: 32 low nibbles (4 bits each, packed 2 per byte)
///
/// Reconstruction: val = d * ((nibble | (high_bit << 4)) - 16)
fn dequant_q5_0(data: &[u8], n_elements: usize) -> Result<Vec<f32>, QuantError> {
    let block_size = 22;
    let elements_per_block = 32;
    let n_blocks = n_elements.div_ceil(elements_per_block);
    if data.len() < n_blocks * block_size {
        return Err(QuantError::InvalidBlockSize {
            expected: n_blocks * block_size,
            got: data.len(),
        });
    }
    let mut out = Vec::with_capacity(n_elements);
    for block_idx in 0..n_blocks {
        let b = block_idx * block_size;
        let d = half::f16::from_le_bytes([data[b], data[b + 1]]).to_f32();
        let qh = u32::from_le_bytes([data[b + 2], data[b + 3], data[b + 4], data[b + 5]]);
        let qs = &data[b + 6..b + 22];

        for j in 0..elements_per_block {
            if out.len() >= n_elements {
                break;
            }
            let nibble = if j % 2 == 0 {
                qs[j / 2] & 0x0F
            } else {
                qs[j / 2] >> 4
            };
            let high_bit = ((qh >> j) & 1) as u8;
            let val = (nibble | (high_bit << 4)) as i32 - 16;
            out.push(d * val as f32);
        }
    }
    out.truncate(n_elements);
    Ok(out)
}

/// Q6_K: 256 elements per block, 210 bytes (super-block with 16 sub-blocks of 16).
///
/// Block layout (210 bytes):
///   - 128 bytes: low 4 bits of quantized values (ql), packed 2 per byte
///   - 64 bytes: high 2 bits of quantized values (qh), packed 4 per byte
///   - 16 bytes: i8 scales for 16 sub-blocks
///   - 2 bytes: f16 super-scale (d)
///
/// Reconstruction: val = d * scale[sub] * (q6_val - 32)
/// where q6_val = low_4bits | (high_2bits << 4)
fn dequant_q6_k(data: &[u8], n_elements: usize) -> Result<Vec<f32>, QuantError> {
    let block_size = 210;
    let elements_per_block = 256;
    let n_blocks = n_elements.div_ceil(elements_per_block);
    if data.len() < n_blocks * block_size {
        return Err(QuantError::InvalidBlockSize {
            expected: n_blocks * block_size,
            got: data.len(),
        });
    }
    let mut out = Vec::with_capacity(n_elements);
    for block_idx in 0..n_blocks {
        let b = block_idx * block_size;
        let ql = &data[b..b + 128];       // low 4 bits
        let qh = &data[b + 128..b + 192]; // high 2 bits
        let scales = &data[b + 192..b + 208]; // i8 scales
        let d = half::f16::from_le_bytes([data[b + 208], data[b + 209]]).to_f32();

        for j in 0..elements_per_block {
            if out.len() >= n_elements {
                break;
            }
            // Low 4 bits: 2 per byte in ql
            let low4 = if j % 2 == 0 {
                ql[j / 2] & 0x0F
            } else {
                ql[j / 2] >> 4
            };
            // High 2 bits: 4 per byte in qh
            let qh_byte = qh[j / 4];
            let high2 = (qh_byte >> ((j % 4) * 2)) & 0x03;
            let q6_val = (low4 | (high2 << 4)) as i32 - 32;
            let sub_idx = j / 16;
            let scale = scales[sub_idx] as i8 as f32;
            out.push(d * scale * q6_val as f32);
        }
    }
    out.truncate(n_elements);
    Ok(out)
}

/// Q4_K_M: 256 elements per block (super-block with eight sub-blocks).
///
/// Block layout (144 bytes total):
///   - 2 bytes: f16 d (super-block scale)
///   - 2 bytes: f16 dmin (super-block min)
///   - 12 bytes: scales/mins for 8 sub-blocks (6-bit packed)
///   - 128 bytes: 256 4-bit quantized values (128 bytes = 256 nibbles)
///
/// CRITICAL: Upper scale bits use << 2 (NOT << 4). This creates a 6-bit value
/// from 2 low bits (from >> 6) + 4 upper bits (from scales_raw[i+8]).
fn dequant_q4_k(data: &[u8], n_elements: usize) -> Result<Vec<f32>, QuantError> {
    let block_size = 144;
    let elements_per_block = 256;
    let n_blocks = n_elements.div_ceil(elements_per_block);
    if data.len() < n_blocks * block_size {
        return Err(QuantError::InvalidBlockSize {
            expected: n_blocks * block_size,
            got: data.len(),
        });
    }
    let mut out = Vec::with_capacity(n_elements);
    for block_idx in 0..n_blocks {
        let b = block_idx * block_size;
        let d = half::f16::from_le_bytes([data[b], data[b + 1]]).to_f32();
        let dmin = half::f16::from_le_bytes([data[b + 2], data[b + 3]]).to_f32();
        let scales_raw = &data[b + 4..b + 16];

        // Decode 8 scale/min pairs from 12 bytes (6-bit packing)
        let mut sc = [0u8; 8];
        let mut mn = [0u8; 8];
        for i in 0..4 {
            sc[i] = scales_raw[i] & 0x3F;
            mn[i] = scales_raw[i + 4] & 0x3F;
        }
        for i in 0..4 {
            let upper_sc = (scales_raw[i + 8] & 0x0F) << 2; // CRITICAL: << 2 not << 4
            let upper_mn = (scales_raw[i + 8] >> 4) << 2; // CRITICAL: << 2 not << 4
            sc[4 + i] = (scales_raw[i] >> 6) | upper_sc;
            mn[4 + i] = (scales_raw[i + 4] >> 6) | upper_mn;
        }

        let qs = &data[b + 16..b + 144];

        for sub in 0..8 {
            let sub_d = d * sc[sub] as f32;
            let sub_m = dmin * mn[sub] as f32;
            let qs_offset = sub * 16;
            for j in 0..32 {
                if out.len() >= n_elements {
                    break;
                }
                let byte_idx = qs_offset + j / 2;
                let nibble = if j % 2 == 0 {
                    qs[byte_idx] & 0x0F
                } else {
                    qs[byte_idx] >> 4
                };
                out.push(sub_d * nibble as f32 - sub_m);
            }
        }
    }
    out.truncate(n_elements);
    Ok(out)
}

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

    #[test]
    fn test_dequant_f32_roundtrip() {
        let values: Vec<f32> = vec![1.0, -2.5, std::f32::consts::PI, 0.0];
        let mut bytes = Vec::new();
        for v in &values {
            bytes.extend_from_slice(&v.to_le_bytes());
        }
        let result = dequantize(&bytes, GgmlType::F32, 4).unwrap();
        assert_eq!(result, values);
    }

    #[test]
    fn test_dequant_f16_roundtrip() {
        let values: Vec<f32> = vec![1.0, -2.0, 0.5];
        let mut bytes = Vec::new();
        for v in &values {
            bytes.extend_from_slice(&half::f16::from_f32(*v).to_le_bytes());
        }
        let result = dequantize(&bytes, GgmlType::F16, 3).unwrap();
        for (a, b) in result.iter().zip(values.iter()) {
            assert!((a - b).abs() < 0.01, "f16 roundtrip: {a} vs {b}");
        }
    }

    #[test]
    fn test_dequant_q8_0_basic() {
        let scale = half::f16::from_f32(1.0);
        let mut block = Vec::new();
        block.extend_from_slice(&scale.to_le_bytes());
        for i in 0..32i8 {
            block.push(i as u8);
        }
        assert_eq!(block.len(), 34);
        let result = dequantize(&block, GgmlType::Q8_0, 32).unwrap();
        assert_eq!(result.len(), 32);
        for (i, val) in result.iter().enumerate() {
            assert!(
                (val - i as f32).abs() < 0.01,
                "q8_0[{i}]: expected {i}, got {val}",
            );
        }
    }

    #[test]
    fn test_dequant_q8_0_with_scale() {
        let scale = half::f16::from_f32(0.5);
        let mut block = Vec::new();
        block.extend_from_slice(&scale.to_le_bytes());
        for i in 0..32i8 {
            block.push(i as u8);
        }
        let result = dequantize(&block, GgmlType::Q8_0, 32).unwrap();
        assert!((result[10] - 5.0).abs() < 0.01);
        assert!((result[20] - 10.0).abs() < 0.01);
    }

    #[test]
    fn test_dequant_q4_k_m_reference() {
        let mut block = vec![0u8; 144];
        let d = half::f16::from_f32(2.0);
        let dmin = half::f16::from_f32(1.0);
        block[0..2].copy_from_slice(&d.to_le_bytes());
        block[2..4].copy_from_slice(&dmin.to_le_bytes());
        let scales_raw = &mut block[4..16];
        scales_raw[0] = 3;
        scales_raw[4] = 1;
        let qs = &mut block[16..144];
        qs[0] = 0x52; // low nibble = 2, high nibble = 5

        let result = dequantize(&block, GgmlType::Q4K, 256).unwrap();
        assert_eq!(result.len(), 256);
        // Sub-block 0: sub_d = 2.0 * 3 = 6.0, sub_m = 1.0 * 1 = 1.0
        // Element 0: nibble=2 -> 6.0 * 2 - 1.0 = 11.0
        // Element 1: nibble=5 -> 6.0 * 5 - 1.0 = 29.0
        assert!(
            (result[0] - 11.0).abs() < 0.1,
            "q4_k_m[0]: expected 11.0, got {}",
            result[0]
        );
        assert!(
            (result[1] - 29.0).abs() < 0.1,
            "q4_k_m[1]: expected 29.0, got {}",
            result[1]
        );
        assert!(
            (result[32] - 0.0).abs() < 0.1,
            "q4_k_m[32]: expected 0.0, got {}",
            result[32]
        );
    }

    #[test]
    fn test_dequant_q5_0_basic() {
        // Build a Q5_0 block: 22 bytes = 2 (f16 scale) + 4 (high bits) + 16 (nibbles)
        let mut block = vec![0u8; 22];
        let scale = half::f16::from_f32(1.0);
        block[0..2].copy_from_slice(&scale.to_le_bytes());
        // qh = 0 (no high bits set)
        // qs[0] = 0x21 -> nibble[0]=1, nibble[1]=2
        block[6] = 0x21;
        let result = dequantize(&block, GgmlType::Q5_0, 32).unwrap();
        assert_eq!(result.len(), 32);
        // val = d * ((nibble | 0) - 16) = 1.0 * (1 - 16) = -15.0
        assert!((result[0] - (-15.0)).abs() < 0.01, "q5_0[0]={}", result[0]);
        // val = 1.0 * (2 - 16) = -14.0
        assert!((result[1] - (-14.0)).abs() < 0.01, "q5_0[1]={}", result[1]);
    }

    #[test]
    fn test_dequant_q5_0_with_high_bit() {
        let mut block = vec![0u8; 22];
        let scale = half::f16::from_f32(1.0);
        block[0..2].copy_from_slice(&scale.to_le_bytes());
        // Set high bit for element 0: qh bit 0 = 1
        block[2] = 0x01; // qh LE byte 0
        // qs[0] low nibble = 5
        block[6] = 0x05;
        let result = dequantize(&block, GgmlType::Q5_0, 32).unwrap();
        // val = 1.0 * ((5 | (1<<4)) - 16) = 1.0 * (21 - 16) = 5.0
        assert!((result[0] - 5.0).abs() < 0.01, "q5_0 high bit: {}", result[0]);
    }

    #[test]
    fn test_dequant_q6_k_basic() {
        // Build a Q6_K block: 210 bytes
        let mut block = vec![0u8; 210];
        // d (super-scale) at offset 208
        let d = half::f16::from_f32(1.0);
        block[208..210].copy_from_slice(&d.to_le_bytes());
        // scales[0] = 2 (i8) at offset 192
        block[192] = 2;
        // ql[0] low nibble = 3 -> element 0 low4 = 3
        block[0] = 0x03;
        // qh[0] = 0 -> high2 = 0 for elements 0-3
        let result = dequantize(&block, GgmlType::Q6K, 256).unwrap();
        assert_eq!(result.len(), 256);
        // q6_val = (3 | 0) - 32 = -29, scale=2, d=1.0
        // val = 1.0 * 2.0 * (-29.0) = -58.0
        assert!((result[0] - (-58.0)).abs() < 0.01, "q6_k[0]={}", result[0]);
    }

    #[test]
    fn test_unsupported_type_errors() {
        let result = dequantize(&[], GgmlType::Iq1S, 1);
        assert!(matches!(result, Err(QuantError::UnsupportedType(_))));
    }

    #[test]
    fn test_block_size_known_types() {
        assert_eq!(block_size_bytes(GgmlType::F32), Some(4));
        assert_eq!(block_size_bytes(GgmlType::F16), Some(2));
        assert_eq!(block_size_bytes(GgmlType::Q8_0), Some(34));
        assert_eq!(block_size_bytes(GgmlType::Q5_0), Some(22));
        assert_eq!(block_size_bytes(GgmlType::Q4K), Some(144));
        assert_eq!(block_size_bytes(GgmlType::Q6K), Some(210));
        assert_eq!(block_size_bytes(GgmlType::Iq1S), None);
    }

    #[test]
    fn test_block_elements_known_types() {
        assert_eq!(block_elements(GgmlType::Q8_0), Some(32));
        assert_eq!(block_elements(GgmlType::Q5_0), Some(32));
        assert_eq!(block_elements(GgmlType::Q4K), Some(256));
        assert_eq!(block_elements(GgmlType::Q6K), Some(256));
    }
}
