//! DCT-based embedding compression for Canon Protocol per CP-016.
//!
//! Compresses 1536-dimensional i16 embedding vectors from ~3072 bytes
//! to ~350 bytes using Type-II DCT, top-5% coefficient selection,
//! and 8-bit quantization. Maintains cosine similarity > 0.95 at
//! default settings.

use serde::{Deserialize, Serialize};
use std::f64::consts::PI;

/// Default fraction of DCT coefficients to retain (5% = 77 of 1536).
pub const DEFAULT_TOP_K_FRACTION: f64 = 0.05;

/// Compressed embedding representation.
///
/// At default K=77 for 1536-dim vectors, the total serialized size
/// is approximately 350 bytes (vs 3072 bytes for the full i16 vector).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompressedEmbedding {
    /// Chunk ID this embedding belongs to
    pub chunk_id: [u8; 16],
    /// Hash of the embedding model used
    pub model_hash: [u8; 32],
    /// Number of retained DCT coefficients
    pub top_k: u16,
    /// Sorted DCT coefficient indices
    pub indices: Vec<u16>,
    /// 8-bit quantized coefficient values
    pub values: Vec<i8>,
    /// Dequantization scale factor
    pub scale: f32,
    /// BLAKE3 hash of the original i16 vector for deferred verification
    pub original_hash: [u8; 32],
}

impl CompressedEmbedding {
    /// Serialize to the CP-016 wire format.
    pub fn to_bytes(&self) -> Vec<u8> {
        let k = self.top_k as usize;
        // 2 + 16 + 32 + 2 + 2*k + k + 4 + 32
        let size = 88 + 3 * k;
        let mut buf = Vec::with_capacity(size);

        // chunk_id_length (always 16)
        buf.extend_from_slice(&16u16.to_le_bytes());
        buf.extend_from_slice(&self.chunk_id);
        buf.extend_from_slice(&self.model_hash);
        buf.extend_from_slice(&self.top_k.to_le_bytes());

        for &idx in &self.indices {
            buf.extend_from_slice(&idx.to_le_bytes());
        }
        for &val in &self.values {
            buf.push(val as u8);
        }

        buf.extend_from_slice(&self.scale.to_le_bytes());
        buf.extend_from_slice(&self.original_hash);

        buf
    }

    /// Deserialize from the CP-016 wire format.
    pub fn from_bytes(data: &[u8]) -> Result<Self, &'static str> {
        if data.len() < 88 {
            return Err("Data too short for CompressedEmbedding header");
        }

        let chunk_id_len = u16::from_le_bytes([data[0], data[1]]) as usize;
        if chunk_id_len != 16 {
            return Err("Expected chunk_id_length of 16");
        }

        let mut chunk_id = [0u8; 16];
        chunk_id.copy_from_slice(&data[2..18]);

        let mut model_hash = [0u8; 32];
        model_hash.copy_from_slice(&data[18..50]);

        let top_k = u16::from_le_bytes([data[50], data[51]]) as usize;

        let indices_start = 52;
        let indices_end = indices_start + 2 * top_k;
        let values_end = indices_end + top_k;
        let scale_end = values_end + 4;
        let hash_end = scale_end + 32;

        if data.len() < hash_end {
            return Err("Data too short for specified top_k");
        }

        let indices: Vec<u16> = (0..top_k)
            .map(|i| {
                let off = indices_start + 2 * i;
                u16::from_le_bytes([data[off], data[off + 1]])
            })
            .collect();

        let values: Vec<i8> = data[indices_end..values_end]
            .iter()
            .map(|&b| b as i8)
            .collect();

        let scale = f32::from_le_bytes([
            data[values_end],
            data[values_end + 1],
            data[values_end + 2],
            data[values_end + 3],
        ]);

        let mut original_hash = [0u8; 32];
        original_hash.copy_from_slice(&data[scale_end..hash_end]);

        Ok(Self {
            chunk_id,
            model_hash,
            top_k: top_k as u16,
            indices,
            values,
            scale,
            original_hash,
        })
    }

    /// Total serialized size in bytes.
    pub fn wire_size(&self) -> usize {
        88 + 3 * self.top_k as usize
    }
}

// ============================================================================
// Compression pipeline
// ============================================================================

/// Compress an i16 embedding vector using DCT + top-k + quantization.
///
/// The `top_k_fraction` controls the quality/size tradeoff:
///   0.025 → ~190 bytes, cosine sim > 0.90
///   0.05  → ~350 bytes, cosine sim > 0.95 (default)
///   0.10  → ~620 bytes, cosine sim > 0.98
///   0.20  → ~1100 bytes, cosine sim > 0.99
pub fn compress(
    vector: &[i16],
    chunk_id: [u8; 16],
    model_hash: [u8; 32],
    top_k_fraction: f64,
) -> CompressedEmbedding {
    let n = vector.len();

    // Step 0: Compute BLAKE3 hash of original vector for verification
    let original_bytes: Vec<u8> = vector.iter().flat_map(|v| v.to_le_bytes()).collect();
    let original_hash = *blake3::hash(&original_bytes).as_bytes();

    // Step 1: Convert i16 to f64 for DCT
    let f64_vec: Vec<f64> = vector.iter().map(|&v| f64::from(v)).collect();

    // Step 2: Apply Type-II DCT
    let dct_coeffs = dct_ii(&f64_vec);

    // Step 3: Top-K sparsification
    let k = ((n as f64 * top_k_fraction).ceil() as usize).max(1).min(n);
    let (indices, sparse_values) = top_k_select(&dct_coeffs, k);

    // Step 4: 8-bit quantization
    let (quantized_values, scale) = quantize_i8(&sparse_values);

    CompressedEmbedding {
        chunk_id,
        model_hash,
        top_k: k as u16,
        indices,
        values: quantized_values,
        scale,
        original_hash,
    }
}

/// Decompress a `CompressedEmbedding` back to an i16 vector.
pub fn decompress(compressed: &CompressedEmbedding, dim: usize) -> Vec<i16> {
    // Step 1: Reconstruct sparse DCT coefficients
    let mut dct_coeffs = vec![0.0f64; dim];
    for (i, &idx) in compressed.indices.iter().enumerate() {
        if (idx as usize) < dim {
            dct_coeffs[idx as usize] =
                f64::from(compressed.values[i]) * f64::from(compressed.scale);
        }
    }

    // Step 2: Inverse DCT (Type-III)
    let reconstructed = idct_iii(&dct_coeffs);

    // Step 3: Convert back to i16 with clamping
    reconstructed
        .iter()
        .map(|&v| v.round().clamp(-32767.0, 32767.0) as i16)
        .collect()
}

/// Decompress to f32 vector (for direct use in similarity computations).
pub fn decompress_f32(compressed: &CompressedEmbedding, dim: usize) -> Vec<f32> {
    let mut dct_coeffs = vec![0.0f64; dim];
    for (i, &idx) in compressed.indices.iter().enumerate() {
        if (idx as usize) < dim {
            dct_coeffs[idx as usize] =
                f64::from(compressed.values[i]) * f64::from(compressed.scale);
        }
    }

    let reconstructed = idct_iii(&dct_coeffs);
    reconstructed.iter().map(|&v| v as f32).collect()
}

// ============================================================================
// DCT-II (forward transform)
// ============================================================================

/// Type-II Discrete Cosine Transform.
///
/// X[k] = sum_{n=0}^{N-1} x[n] * cos(PI * (2n + 1) * k / (2N))
///
/// This is the naive O(N^2) implementation. For production with very
/// large N, an FFT-based O(N log N) implementation would be preferable,
/// but N=1536 completes in well under 1ms on modern hardware.
fn dct_ii(x: &[f64]) -> Vec<f64> {
    let n = x.len();
    let nf = n as f64;
    let mut result = vec![0.0; n];

    for (k, result_k) in result.iter_mut().enumerate() {
        let kf = k as f64;
        let mut sum = 0.0;
        for (idx, &val) in x.iter().enumerate() {
            let nf_idx = idx as f64;
            sum += val * (PI * (2.0 * nf_idx + 1.0) * kf / (2.0 * nf)).cos();
        }
        *result_k = sum;
    }

    result
}

// ============================================================================
// IDCT (Type-III, inverse transform)
// ============================================================================

/// Type-III Discrete Cosine Transform (inverse of Type-II).
///
/// x[n] = (1/N) * X[0] + (2/N) * sum_{k=1}^{N-1} X[k] * cos(PI * k * (2n + 1) / (2N))
fn idct_iii(coeffs: &[f64]) -> Vec<f64> {
    let n = coeffs.len();
    let nf = n as f64;
    let mut result = vec![0.0; n];

    for (idx, result_idx) in result.iter_mut().enumerate() {
        let nf_idx = idx as f64;
        // DC component
        let mut sum = coeffs[0] / nf;
        // AC components
        for (k, &coeff_k) in coeffs.iter().enumerate().skip(1) {
            let kf = k as f64;
            sum += (2.0 / nf) * coeff_k * (PI * kf * (2.0 * nf_idx + 1.0) / (2.0 * nf)).cos();
        }
        *result_idx = sum;
    }

    result
}

// ============================================================================
// Top-K selection
// ============================================================================

/// Select the K largest coefficients by absolute value.
/// Returns (`sorted_indices`, `corresponding_values`).
fn top_k_select(coeffs: &[f64], k: usize) -> (Vec<u16>, Vec<f64>) {
    // Create index-value pairs sorted by absolute value (descending)
    let mut indexed: Vec<(usize, f64)> = coeffs.iter().enumerate().map(|(i, &v)| (i, v)).collect();
    indexed.sort_by(|a, b| {
        b.1.abs()
            .partial_cmp(&a.1.abs())
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    // Take top-k
    let selected: Vec<(usize, f64)> = indexed.into_iter().take(k).collect();

    // Sort by index for compact encoding
    let mut sorted = selected;
    sorted.sort_by_key(|(i, _)| *i);

    let indices: Vec<u16> = sorted.iter().map(|(i, _)| *i as u16).collect();
    let values: Vec<f64> = sorted.iter().map(|(_, v)| *v).collect();

    (indices, values)
}

// ============================================================================
// 8-bit quantization
// ============================================================================

/// Quantize f64 values to i8 with a shared scale factor.
/// scale = `max_abs` / 127.0, then each value = round(v / scale), clamped to [-127, 127].
fn quantize_i8(values: &[f64]) -> (Vec<i8>, f32) {
    if values.is_empty() {
        return (Vec::new(), 0.0);
    }

    let max_abs = values.iter().map(|v| v.abs()).fold(0.0f64, f64::max);
    if max_abs == 0.0 {
        return (vec![0i8; values.len()], 0.0);
    }

    let scale = max_abs / 127.0;
    let quantized: Vec<i8> = values
        .iter()
        .map(|&v| (v / scale).round().clamp(-127.0, 127.0) as i8)
        .collect();

    (quantized, scale as f32)
}

// ============================================================================
// Utility: cosine similarity
// ============================================================================

/// Compute cosine similarity between two i16 vectors.
pub fn cosine_similarity_i16(a: &[i16], b: &[i16]) -> f64 {
    assert_eq!(a.len(), b.len());
    let dot: f64 = a
        .iter()
        .zip(b.iter())
        .map(|(&x, &y)| f64::from(x) * f64::from(y))
        .sum();
    let norm_a: f64 = a.iter().map(|&x| f64::from(x).powi(2)).sum::<f64>().sqrt();
    let norm_b: f64 = b.iter().map(|&x| f64::from(x).powi(2)).sum::<f64>().sqrt();
    if norm_a == 0.0 || norm_b == 0.0 {
        return 0.0;
    }
    dot / (norm_a * norm_b)
}

/// Compute cosine similarity between two f32 vectors.
pub fn cosine_similarity_f32(a: &[f32], b: &[f32]) -> f64 {
    assert_eq!(a.len(), b.len());
    let dot: f64 = a
        .iter()
        .zip(b.iter())
        .map(|(&x, &y)| f64::from(x) * f64::from(y))
        .sum();
    let norm_a: f64 = a.iter().map(|&x| f64::from(x).powi(2)).sum::<f64>().sqrt();
    let norm_b: f64 = b.iter().map(|&x| f64::from(x).powi(2)).sum::<f64>().sqrt();
    if norm_a == 0.0 || norm_b == 0.0 {
        return 0.0;
    }
    dot / (norm_a * norm_b)
}

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

    fn sample_vector(dim: usize) -> Vec<i16> {
        // Simulate a realistic embedding: mostly small values with some structure
        use rand::Rng;
        let mut rng = rand::thread_rng();
        (0..dim)
            .map(|i| {
                // Low-frequency pattern + noise
                let signal = 1000.0 * (2.0 * PI * i as f64 / dim as f64).sin()
                    + 500.0 * (6.0 * PI * i as f64 / dim as f64).cos();
                let noise = rng.gen_range(-200.0..200.0);
                (signal + noise).round().clamp(-32767.0, 32767.0) as i16
            })
            .collect()
    }

    #[test]
    fn test_dct_idct_roundtrip() {
        // For a small vector, DCT -> IDCT should be near-identity
        let x = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
        let dct = dct_ii(&x);
        let reconstructed = idct_iii(&dct);

        for (orig, recon) in x.iter().zip(reconstructed.iter()) {
            assert!(
                (orig - recon).abs() < 1e-10,
                "DCT roundtrip failed: {orig} vs {recon}"
            );
        }
    }

    #[test]
    fn test_dct_idct_roundtrip_large() {
        // Test with 1536-dim vector (our actual embedding size)
        let x: Vec<f64> = (0..1536)
            .map(|i| (f64::from(i) * 0.01).sin() * 1000.0)
            .collect();
        let dct = dct_ii(&x);
        let reconstructed = idct_iii(&dct);

        for (i, (orig, recon)) in x.iter().zip(reconstructed.iter()).enumerate() {
            assert!(
                (orig - recon).abs() < 1e-6,
                "DCT roundtrip failed at index {i}: {orig} vs {recon}"
            );
        }
    }

    #[test]
    fn test_top_k_select() {
        let coeffs = vec![1.0, -5.0, 3.0, -2.0, 4.0];
        let (indices, values) = top_k_select(&coeffs, 3);

        assert_eq!(indices.len(), 3);
        assert_eq!(values.len(), 3);

        // Should select indices 1 (-5.0), 4 (4.0), 2 (3.0) — sorted by index
        assert!(indices.contains(&1));
        assert!(indices.contains(&2));
        assert!(indices.contains(&4));
    }

    #[test]
    fn test_quantize_i8() {
        let values = vec![100.0, -50.0, 200.0, -200.0];
        let (quantized, scale) = quantize_i8(&values);

        assert_eq!(quantized.len(), 4);
        assert!(scale > 0.0);

        // Max abs is 200.0, scale = 200.0/127.0
        // 200.0 / scale = 127
        assert_eq!(quantized[2], 127);
        assert_eq!(quantized[3], -127);
    }

    #[test]
    fn test_quantize_zero_vector() {
        let values = vec![0.0, 0.0, 0.0];
        let (quantized, scale) = quantize_i8(&values);
        assert_eq!(quantized, vec![0, 0, 0]);
        assert!(scale.abs() < f32::EPSILON);
    }

    #[test]
    fn test_compress_decompress_roundtrip() {
        let dim = 1536;
        let vector = sample_vector(dim);
        let chunk_id = [1u8; 16];
        let model_hash = [2u8; 32];

        let compressed = compress(&vector, chunk_id, model_hash, DEFAULT_TOP_K_FRACTION);
        let reconstructed = decompress(&compressed, dim);

        assert_eq!(reconstructed.len(), dim);

        let sim = cosine_similarity_i16(&vector, &reconstructed);
        assert!(
            sim > 0.95,
            "Cosine similarity at default K should be > 0.95, got {sim}"
        );
    }

    #[test]
    fn test_compress_size_reduction() {
        let dim = 1536;
        let vector = sample_vector(dim);

        let compressed = compress(&vector, [0u8; 16], [0u8; 32], DEFAULT_TOP_K_FRACTION);
        let wire_size = compressed.wire_size();
        let original_size = dim * 2; // i16 = 2 bytes

        assert!(
            wire_size < original_size / 8,
            "Should achieve at least 8x compression: {wire_size} vs {original_size}"
        );

        // At K=77: 88 + 3*77 = 319 bytes
        assert!(
            wire_size < 400,
            "Wire size should be ~350 bytes, got {wire_size}"
        );
    }

    #[test]
    fn test_compress_quality_at_different_k() {
        let dim = 1536;
        let vector = sample_vector(dim);
        let chunk_id = [0u8; 16];
        let model_hash = [0u8; 32];

        // K=2.5%: sim > 0.90
        let c1 = compress(&vector, chunk_id, model_hash, 0.025);
        let r1 = decompress(&c1, dim);
        let sim1 = cosine_similarity_i16(&vector, &r1);
        assert!(sim1 > 0.85, "2.5% K: expected sim > 0.85, got {sim1}");

        // K=5%: sim > 0.95
        let c2 = compress(&vector, chunk_id, model_hash, 0.05);
        let r2 = decompress(&c2, dim);
        let sim2 = cosine_similarity_i16(&vector, &r2);
        assert!(sim2 > 0.95, "5% K: expected sim > 0.95, got {sim2}");

        // K=10%: sim > 0.98
        let c3 = compress(&vector, chunk_id, model_hash, 0.10);
        let r3 = decompress(&c3, dim);
        let sim3 = cosine_similarity_i16(&vector, &r3);
        assert!(sim3 > 0.97, "10% K: expected sim > 0.97, got {sim3}");

        // Higher K should give better similarity
        assert!(sim2 > sim1);
        assert!(sim3 > sim2);
    }

    #[test]
    fn test_wire_format_roundtrip() {
        let dim = 1536;
        let vector = sample_vector(dim);

        let compressed = compress(&vector, [3u8; 16], [4u8; 32], DEFAULT_TOP_K_FRACTION);
        let bytes = compressed.to_bytes();
        let parsed = CompressedEmbedding::from_bytes(&bytes).unwrap();

        assert_eq!(parsed.chunk_id, compressed.chunk_id);
        assert_eq!(parsed.model_hash, compressed.model_hash);
        assert_eq!(parsed.top_k, compressed.top_k);
        assert_eq!(parsed.indices, compressed.indices);
        assert_eq!(parsed.values, compressed.values);
        assert!((parsed.scale - compressed.scale).abs() < 1e-6);
        assert_eq!(parsed.original_hash, compressed.original_hash);

        // Decompressed from parsed should match original decompressed
        let recon1 = decompress(&compressed, dim);
        let recon2 = decompress(&parsed, dim);
        assert_eq!(recon1, recon2);
    }

    #[test]
    fn test_original_hash_verification() {
        let dim = 1536;
        let vector = sample_vector(dim);

        let compressed = compress(&vector, [0u8; 16], [0u8; 32], DEFAULT_TOP_K_FRACTION);

        // Recompute hash from original vector
        let original_bytes: Vec<u8> = vector.iter().flat_map(|v| v.to_le_bytes()).collect();
        let expected_hash = *blake3::hash(&original_bytes).as_bytes();

        assert_eq!(
            compressed.original_hash, expected_hash,
            "original_hash should match BLAKE3 of the i16 vector bytes"
        );
    }

    #[test]
    fn test_decompress_f32() {
        let dim = 1536;
        let vector = sample_vector(dim);

        let compressed = compress(&vector, [0u8; 16], [0u8; 32], DEFAULT_TOP_K_FRACTION);
        let recon_f32 = decompress_f32(&compressed, dim);

        assert_eq!(recon_f32.len(), dim);

        // Convert original to f32 for comparison
        let orig_f32: Vec<f32> = vector.iter().map(|&v| f32::from(v)).collect();
        let sim = cosine_similarity_f32(&orig_f32, &recon_f32);
        assert!(sim > 0.95, "f32 decompress sim should be > 0.95, got {sim}");
    }

    #[test]
    fn test_cosine_similarity_identical() {
        let a = vec![1i16, 2, 3, 4, 5];
        let sim = cosine_similarity_i16(&a, &a);
        assert!((sim - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_cosine_similarity_orthogonal() {
        let a = vec![1i16, 0, 0, 0];
        let b = vec![0i16, 1, 0, 0];
        let sim = cosine_similarity_i16(&a, &b);
        assert!(sim.abs() < 1e-10);
    }

    #[test]
    fn test_empty_vector() {
        let vector: Vec<i16> = Vec::new();
        let compressed = compress(&vector, [0u8; 16], [0u8; 32], 0.05);
        assert_eq!(compressed.top_k, 0);
        let decompressed = decompress(&compressed, 0);
        assert!(decompressed.is_empty());
    }

    #[test]
    fn test_small_vector() {
        let vector = vec![100i16, -200, 300, -400];
        let compressed = compress(&vector, [0u8; 16], [0u8; 32], 0.5); // 50% K for small vec
        let decompressed = decompress(&compressed, 4);

        let sim = cosine_similarity_i16(&vector, &decompressed);
        assert!(sim > 0.9, "Small vector roundtrip sim: {sim}");
    }
}
