constraint-theory-core 2.2.0

//! Pythagorean Quantizer - Unified Quantization with Constraint Preservation
//!
//! This module provides the `PythagoreanQuantizer` which synthesizes multiple
//! quantization technologies for exact constraint satisfaction:
//!
//! - **TurboQuant**: Near-optimal distortion, works online, O(d log d)
//! - **BitNet**: Ternary weights {-1, 0, 1} for LLM inference
//! - **PolarQuant**: Exact unit norm preservation via polar coordinate quantization
//!
//! # Architecture
//!
//! ```text
//! Input ──► [Mode Selector] ──► [Quantizer] ──► [Constraint Layer]
//!
//! Modes:
//! • TERNARY  (BitNet): {-1, 0, 1} for LLM weights
//! • POLAR    (PolarQuant): Exact unit norm preservation
//! • TURBO    (TurboQuant): Near-optimal distortion
//! • HYBRID: Auto-select based on input characteristics
//! ```
//!
//! # Example
//!
//! ```
//! use constraint_theory_core::quantizer::{PythagoreanQuantizer, QuantizationMode};
//!
//! // Create a quantizer with POLAR mode for unit norm preservation
//! let quantizer = PythagoreanQuantizer::new(QuantizationMode::Polar, 8);
//!
//! // Quantize a vector
//! let vector = vec![0.6, 0.8, 0.0, 0.0];
//! let result = quantizer.quantize(&vector);
//!
//! // Result preserves unit norm exactly
//! let norm: f64 = result.data.iter().map(|x| x * x).sum::<f64>().sqrt();
//! assert!((norm - 1.0).abs() < 0.01);
//! ```

use std::f64;

/// Quantization modes supported by PythagoreanQuantizer.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum QuantizationMode {
    /// Ternary quantization (BitNet style): {-1, 0, 1}
    /// Best for: LLM weights, sparse representations
    Ternary,
    
    /// Polar coordinate quantization (PolarQuant style)
    /// Best for: Unit norm preservation, embeddings
    Polar,
    
    /// Near-optimal distortion quantization (TurboQuant style)
    /// Best for: Vector databases, general purpose
    Turbo,
    
    /// Auto-select mode based on input characteristics
    Hybrid,
}

/// Result of quantization operation.
#[derive(Clone, Debug)]
pub struct QuantizationResult {
    /// Quantized data
    pub data: Vec<f64>,
    /// Quantization mode used
    pub mode: QuantizationMode,
    /// Bits per element
    pub bits: u8,
    /// Mean squared error from original
    pub mse: f64,
    /// Whether constraints are satisfied
    pub constraints_satisfied: bool,
    /// Unit norm preserved (for Polar mode)
    pub unit_norm_preserved: bool,
}

impl QuantizationResult {
    /// Create a new quantization result.
    pub fn new(data: Vec<f64>, mode: QuantizationMode, bits: u8) -> Self {
        Self {
            data,
            mode,
            bits,
            mse: 0.0,
            constraints_satisfied: true,
            unit_norm_preserved: true,
        }
    }
    
    /// Compute the norm of the quantized vector.
    pub fn norm(&self) -> f64 {
        self.data.iter().map(|x| x * x).sum::<f64>().sqrt()
    }
    
    /// Check if unit norm is preserved within tolerance.
    pub fn check_unit_norm(&self, tolerance: f64) -> bool {
        (self.norm() - 1.0).abs() < tolerance
    }
}

/// Pythagorean Quantizer - Unified quantization with constraint preservation.
///
/// Integrates TurboQuant, BitNet, and PolarQuant technologies with
/// Pythagorean snapping for exact constraint satisfaction.
#[derive(Clone, Debug)]
pub struct PythagoreanQuantizer {
    /// Quantization mode
    pub mode: QuantizationMode,
    /// Bits per element (1 for ternary, 4-8 for others)
    pub bits: u8,
    /// Maximum denominator for Pythagorean ratios
    max_denominator: usize,
}

impl PythagoreanQuantizer {
    /// Create a new Pythagorean quantizer.
    ///
    /// # Arguments
    ///
    /// * `mode` - Quantization mode to use
    /// * `bits` - Bits per element (1 for ternary, 4-8 for others)
    ///
    /// # Example
    ///
    /// ```
    /// use constraint_theory_core::quantizer::{PythagoreanQuantizer, QuantizationMode};
    ///
    /// let quantizer = PythagoreanQuantizer::new(QuantizationMode::Ternary, 1);
    /// ```
    pub fn new(mode: QuantizationMode, bits: u8) -> Self {
        Self {
            mode,
            bits: bits.max(1),
            max_denominator: 100,
        }
    }
    
    /// Create a quantizer optimized for LLM weights (ternary).
    pub fn for_llm() -> Self {
        Self::new(QuantizationMode::Ternary, 1)
    }
    
    /// Create a quantizer optimized for embeddings (polar).
    pub fn for_embeddings() -> Self {
        Self::new(QuantizationMode::Polar, 8)
    }
    
    /// Create a quantizer optimized for vector databases (turbo).
    pub fn for_vector_db() -> Self {
        Self::new(QuantizationMode::Turbo, 4)
    }
    
    /// Create a hybrid quantizer that auto-selects mode.
    pub fn hybrid() -> Self {
        Self::new(QuantizationMode::Hybrid, 4)
    }
    
    /// Quantize data with constraint preservation.
    ///
    /// # Arguments
    ///
    /// * `data` - Input data to quantize
    ///
    /// # Returns
    ///
    /// Quantization result with preserved constraints
    ///
    /// # Example
    ///
    /// ```
    /// use constraint_theory_core::quantizer::{PythagoreanQuantizer, QuantizationMode};
    ///
    /// let quantizer = PythagoreanQuantizer::new(QuantizationMode::Polar, 8);
    /// let data = vec![0.6, 0.8, 0.0, 0.0];
    /// let result = quantizer.quantize(&data);
    ///
    /// assert_eq!(result.data.len(), 4);
    /// ```
    pub fn quantize(&self, data: &[f64]) -> QuantizationResult {
        let mode = self.select_mode(data);
        
        let (quantized, mse) = match mode {
            QuantizationMode::Ternary => self.quantize_ternary(data),
            QuantizationMode::Polar => self.quantize_polar(data),
            QuantizationMode::Turbo => self.quantize_turbo(data),
            QuantizationMode::Hybrid => self.quantize_hybrid(data),
        };
        
        let mut result = QuantizationResult::new(quantized, mode, self.bits);
        result.mse = mse;
        result.unit_norm_preserved = self.check_unit_norm(&result.data);
        result.constraints_satisfied = result.unit_norm_preserved || mode != QuantizationMode::Polar;
        
        result
    }
    
    /// Auto-select quantization mode based on input characteristics.
    fn select_mode(&self, data: &[f64]) -> QuantizationMode {
        if self.mode != QuantizationMode::Hybrid {
            return self.mode;
        }
        
        // Check if input is already unit normalized
        let norm: f64 = data.iter().map(|x| x * x).sum::<f64>().sqrt();
        let is_unit_norm = (norm - 1.0).abs() < 0.01;
        
        // Check sparsity (for ternary mode)
        let threshold = 0.1;
        let sparse_count = data.iter().filter(|&&x| x.abs() < threshold).count();
        let sparsity = sparse_count as f64 / data.len() as f64;
        
        if is_unit_norm {
            QuantizationMode::Polar
        } else if sparsity > 0.5 {
            QuantizationMode::Ternary
        } else {
            QuantizationMode::Turbo
        }
    }
    
    /// Ternary quantization (BitNet style): {-1, 0, 1}.
    ///
    /// Achieves 16x memory reduction for LLM weights.
    fn quantize_ternary(&self, data: &[f64]) -> (Vec<f64>, f64) {
        // Compute threshold for zero bucket
        let mean_abs: f64 = data.iter().map(|x| x.abs()).sum::<f64>() / data.len().max(1) as f64;
        let threshold = mean_abs * 0.1; // Small values -> 0
        
        let quantized: Vec<f64> = data.iter().map(|&x| {
            if x.abs() < threshold {
                0.0
            } else if x > 0.0 {
                1.0
            } else {
                -1.0
            }
        }).collect();
        
        let mse: f64 = data.iter()
            .zip(quantized.iter())
            .map(|(o, q)| (o - q).powi(2))
            .sum::<f64>() / data.len().max(1) as f64;
        
        (quantized, mse)
    }
    
    /// Polar coordinate quantization (PolarQuant style).
    ///
    /// Preserves unit norm exactly via polar coordinate quantization.
    fn quantize_polar(&self, data: &[f64]) -> (Vec<f64>, f64) {
        let n = data.len();
        if n < 2 {
            return (data.to_vec(), 0.0);
        }
        
        // Compute current norm
        let norm: f64 = data.iter().map(|x| x * x).sum::<f64>().sqrt();
        if norm < 1e-10 {
            return (vec![1.0], 0.0);
        }
        
        // Normalize first
        let normalized: Vec<f64> = data.iter().map(|&x| x / norm).collect();
        
        // Convert to polar coordinates for each pair
        let mut quantized = vec![0.0; n];
        
        for i in (0..n).step_by(2) {
            if i + 1 < n {
                let (q0, q1) = self.quantize_polar_pair(normalized[i], normalized[i + 1]);
                quantized[i] = q0;
                quantized[i + 1] = q1;
            } else {
                // Odd dimension - snap to nearest Pythagorean ratio
                quantized[i] = self.snap_to_pythagorean(normalized[i]);
            }
        }
        
        // Re-normalize to ensure exact unit norm
        let q_norm: f64 = quantized.iter().map(|x| x * x).sum::<f64>().sqrt();
        if q_norm > 1e-10 {
            quantized = quantized.iter().map(|&x| x / q_norm).collect();
        }
        
        let mse: f64 = normalized.iter()
            .zip(quantized.iter())
            .map(|(o, q)| (o - q).powi(2))
            .sum::<f64>() / n as f64;
        
        (quantized, mse)
    }
    
    /// Quantize a 2D point using polar coordinates.
    fn quantize_polar_pair(&self, x: f64, y: f64) -> (f64, f64) {
        // Convert to angle
        let angle = y.atan2(x);
        
        // Snap angle to nearest Pythagorean angle
        let snapped_angle = self.snap_angle_to_pythagorean(angle);
        
        // Convert back to Cartesian (unit norm preserved)
        (snapped_angle.cos(), snapped_angle.sin())
    }
    
    /// Snap an angle to the nearest Pythagorean angle.
    fn snap_angle_to_pythagorean(&self, angle: f64) -> f64 {
        // Angles corresponding to common Pythagorean triples
        let pythagorean_angles: &[f64] = &[
            0.0, std::f64::consts::FRAC_PI_2, std::f64::consts::PI, -std::f64::consts::FRAC_PI_2,
            // 3-4-5 triangle: atan(4/3) ≈ 0.927 radians
            (4.0_f64 / 3.0).atan(),
            (3.0_f64 / 4.0).atan(),
            // 5-12-13 triangle
            (12.0_f64 / 5.0).atan(),
            (5.0_f64 / 12.0).atan(),
            // 8-15-17 triangle
            (15.0_f64 / 8.0).atan(),
            (8.0_f64 / 15.0).atan(),
            // 45 degrees
            std::f64::consts::FRAC_PI_4,
            // 30 degrees
            std::f64::consts::FRAC_PI_6,
            // 60 degrees
            std::f64::consts::FRAC_PI_3,
        ];
        
        let mut best = angle;
        let mut min_diff = f64::MAX;
        
        for &pyth_angle in pythagorean_angles {
            // Handle angle wrapping
            let diff = ((angle - pyth_angle).abs() % std::f64::consts::TAU)
                .min((pyth_angle - angle).abs() % std::f64::consts::TAU);
            if diff < min_diff {
                min_diff = diff;
                best = pyth_angle;
            }
        }
        
        best
    }
    
    /// Turbo quantization (TurboQuant style).
    ///
    /// Near-optimal distortion: D(b,d) ≤ 2.7 · D*(b,d)
    fn quantize_turbo(&self, data: &[f64]) -> (Vec<f64>, f64) {
        let n = data.len();
        if n == 0 {
            return (vec![], 0.0);
        }
        
        // Compute statistics
        let min_val = data.iter().cloned().fold(f64::INFINITY, f64::min);
        let max_val = data.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
        let range = max_val - min_val;
        
        if range < 1e-10 {
            return (vec![min_val; n], 0.0);
        }
        
        // Number of quantization levels
        let levels = (1 << self.bits) as f64; // 2^bits
        
        // Quantize each value
        let quantized: Vec<f64> = data.iter().map(|&x| {
            // Scale to [0, levels-1]
            let scaled = ((x - min_val) / range * (levels - 1.0)).round();
            // Snap to Pythagorean ratio if close
            let snapped = self.snap_to_pythagorean(scaled / (levels - 1.0));
            // Scale back
            min_val + snapped * range
        }).collect();
        
        let mse: f64 = data.iter()
            .zip(quantized.iter())
            .map(|(o, q)| (o - q).powi(2))
            .sum::<f64>() / n as f64;
        
        (quantized, mse)
    }
    
    /// Hybrid quantization - combines best aspects of all modes.
    fn quantize_hybrid(&self, data: &[f64]) -> (Vec<f64>, f64) {
        let mode = self.select_mode(data);
        match mode {
            QuantizationMode::Ternary => self.quantize_ternary(data),
            QuantizationMode::Polar => self.quantize_polar(data),
            QuantizationMode::Turbo => self.quantize_turbo(data),
            QuantizationMode::Hybrid => self.quantize_turbo(data), // Default to Turbo
        }
    }
    
    /// Snap a value to the nearest Pythagorean ratio.
    ///
    /// Pythagorean ratios are of the form a/c or b/c where a² + b² = c².
    pub fn snap_to_pythagorean(&self, value: f64) -> f64 {
        // Common Pythagorean ratios from primitive triples
        let pythagorean_ratios: &[f64] = &[
            0.0, 1.0,
            3.0/5.0, 4.0/5.0,
            5.0/13.0, 12.0/13.0,
            8.0/17.0, 15.0/17.0,
            7.0/25.0, 24.0/25.0,
            20.0/29.0, 21.0/29.0,
            9.0/41.0, 40.0/41.0,
            0.5, 0.7071067811865476, // sqrt(2)/2
        ];
        
        let mut best = value;
        let mut min_dist = f64::MAX;
        
        for &ratio in pythagorean_ratios {
            let dist = (value - ratio).abs();
            if dist < min_dist {
                min_dist = dist;
                best = ratio;
            }
        }
        
        best
    }
    
    /// Snap to Pythagorean lattice with explicit rational representation.
    ///
    /// # Arguments
    ///
    /// * `value` - Value to snap
    /// * `max_denominator` - Maximum denominator for rational approximation
    ///
    /// # Returns
    ///
    /// Tuple of (snapped_value, numerator, denominator)
    pub fn snap_to_lattice(&self, value: f64, max_denominator: usize) -> (f64, i64, u64) {
        // Find best rational approximation with Pythagorean constraint
        let mut best_val = value;
        let mut best_num = value.round() as i64;
        let mut best_den = 1u64;
        let mut best_err = f64::MAX;
        
        // Check Pythagorean triples up to max_denominator
        for c in 2..=max_denominator {
            for a in 1..c {
                let b_sq = (c * c - a * a) as f64;
                if b_sq > 0.0 {
                    let b = b_sq.sqrt() as usize;
                    if b * b == (c * c - a * a) {
                        // This is a valid Pythagorean triple
                        let ratio_a = a as f64 / c as f64;
                        let ratio_b = b as f64 / c as f64;
                        
                        let err_a = (value - ratio_a).abs();
                        if err_a < best_err {
                            best_err = err_a;
                            best_val = ratio_a;
                            best_num = a as i64;
                            best_den = c as u64;
                        }
                        
                        let err_b = (value - ratio_b).abs();
                        if err_b < best_err {
                            best_err = err_b;
                            best_val = ratio_b;
                            best_num = b as i64;
                            best_den = c as u64;
                        }
                    }
                }
            }
        }
        
        (best_val, best_num, best_den)
    }
    
    /// Check if unit norm is preserved within tolerance.
    fn check_unit_norm(&self, data: &[f64]) -> bool {
        let norm: f64 = data.iter().map(|x| x * x).sum::<f64>().sqrt();
        (norm - 1.0).abs() < 0.01
    }
    
    /// Batch quantization for multiple vectors.
    ///
    /// # Arguments
    ///
    /// * `vectors` - Slice of vectors to quantize
    ///
    /// # Returns
    ///
    /// Vector of quantization results
    pub fn quantize_batch(&self, vectors: &[Vec<f64>]) -> Vec<QuantizationResult> {
        vectors.iter().map(|v| self.quantize(v)).collect()
    }
}

impl Default for PythagoreanQuantizer {
    fn default() -> Self {
        Self::hybrid()
    }
}

/// A rational number for exact representation.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Rational {
    /// Numerator
    pub num: i64,
    /// Denominator (always positive)
    pub den: u64,
}

impl Rational {
    /// Create a new rational number.
    pub fn new(num: i64, den: u64) -> Self {
        Self { num, den }
    }
    
    /// Convert to floating point.
    pub fn to_f64(&self) -> f64 {
        self.num as f64 / self.den as f64
    }
    
    /// Check if this is a Pythagorean ratio (part of a Pythagorean triple).
    pub fn is_pythagorean(&self) -> bool {
        // Check if numerator² + something² = denominator²
        let a = self.num.unsigned_abs() as u64;
        let c = self.den;
        
        if c == 0 {
            return false;
        }
        
        if a > c {
            return false;
        }
        
        let b_sq = c * c - a * a;
        let b = (b_sq as f64).sqrt() as u64;
        b * b == b_sq
    }
}

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

    #[test]
    fn test_quantization_modes() {
        let data = vec![0.6, 0.8, 0.0, 0.0];
        
        // Test Ternary mode
        let q = PythagoreanQuantizer::new(QuantizationMode::Ternary, 1);
        let result = q.quantize(&data);
        assert_eq!(result.mode, QuantizationMode::Ternary);
        
        // Test Polar mode
        let q = PythagoreanQuantizer::new(QuantizationMode::Polar, 8);
        let result = q.quantize(&data);
        assert!(result.check_unit_norm(0.1));
        
        // Test Turbo mode
        let q = PythagoreanQuantizer::new(QuantizationMode::Turbo, 4);
        let result = q.quantize(&data);
        assert_eq!(result.mode, QuantizationMode::Turbo);
    }

    #[test]
    fn test_polar_unit_norm() {
        let q = PythagoreanQuantizer::for_embeddings();
        
        // Test with various unit vectors
        let vectors = vec![
            vec![1.0, 0.0, 0.0, 0.0],
            vec![0.707, 0.707, 0.0, 0.0],
            vec![0.6, 0.8, 0.0, 0.0],
            vec![0.5, 0.5, 0.5, 0.5],
        ];
        
        for v in vectors {
            let result = q.quantize(&v);
            assert!(result.check_unit_norm(0.1), "Failed for vector {:?}", v);
        }
    }

    #[test]
    fn test_ternary_quantization() {
        let q = PythagoreanQuantizer::for_llm();
        let data = vec![-0.8, -0.1, 0.1, 0.9];
        let result = q.quantize(&data);
        
        // All values should be -1, 0, or 1
        for &val in &result.data {
            assert!(val == -1.0 || val == 0.0 || val == 1.0);
        }
    }

    #[test]
    fn test_snap_to_pythagorean() {
        let q = PythagoreanQuantizer::new(QuantizationMode::Polar, 8);
        
        // 0.6 should snap to 3/5
        let snapped = q.snap_to_pythagorean(0.6);
        assert!((snapped - 0.6).abs() < 0.01);
        
        // 0.8 should snap to 4/5
        let snapped = q.snap_to_pythagorean(0.8);
        assert!((snapped - 0.8).abs() < 0.01);
    }

    #[test]
    fn test_snap_to_lattice() {
        let q = PythagoreanQuantizer::new(QuantizationMode::Polar, 8);
        
        let (val, num, den) = q.snap_to_lattice(0.6, 20);
        assert_eq!(num, 3);
        assert_eq!(den, 5);
        assert!((val - 0.6).abs() < 0.01);
    }

    #[test]
    fn test_hybrid_mode_selection() {
        let q = PythagoreanQuantizer::hybrid();
        
        // Unit norm vector -> should select Polar
        let unit = vec![0.6, 0.8];
        assert_eq!(q.select_mode(&unit), QuantizationMode::Polar);
        
        // Sparse vector -> should select Ternary
        let sparse = vec![0.01, 0.02, 0.0, 0.0, 0.0, 0.0];
        assert_eq!(q.select_mode(&sparse), QuantizationMode::Ternary);
        
        // Dense vector -> should select Turbo
        let dense = vec![0.5, 0.6, 0.7, 0.8];
        assert_eq!(q.select_mode(&dense), QuantizationMode::Turbo);
    }

    #[test]
    fn test_rational() {
        let r = Rational::new(3, 5);
        assert!((r.to_f64() - 0.6).abs() < 1e-10);
        assert!(r.is_pythagorean());
        
        let r = Rational::new(4, 5);
        assert!((r.to_f64() - 0.8).abs() < 1e-10);
        assert!(r.is_pythagorean());
        
        let r = Rational::new(1, 3);
        assert!(!r.is_pythagorean());
    }

    #[test]
    fn test_batch_quantization() {
        let q = PythagoreanQuantizer::for_embeddings();
        let vectors = vec![
            vec![0.6, 0.8],
            vec![1.0, 0.0],
            vec![0.707, 0.707],
        ];
        
        let results = q.quantize_batch(&vectors);
        assert_eq!(results.len(), 3);
        
        for result in results {
            assert!(result.check_unit_norm(0.1));
        }
    }

    #[test]
    fn test_empty_input() {
        let q = PythagoreanQuantizer::hybrid();
        let result = q.quantize(&[]);
        assert!(result.data.is_empty());
    }

    #[test]
    fn test_single_element() {
        let q = PythagoreanQuantizer::new(QuantizationMode::Polar, 8);
        let result = q.quantize(&[1.0]);
        assert_eq!(result.data.len(), 1);
    }
}