sketch_oxide 0.1.6

//! Learned Bloom Filter - ML-Enhanced Membership Testing
//!
//! **EXPERIMENTAL FEATURE** - Use with caution in production systems.
//!
//! A Learned Bloom Filter uses machine learning to predict set membership,
//! achieving 70-80% memory reduction compared to standard Bloom filters.
//!
//! # How It Works
//!
//! 1. **Feature Extraction**: Extract features from keys (hash patterns, bit distributions)
//! 2. **ML Model**: Simple linear model (logistic regression) learns key patterns
//! 3. **Backup Filter**: Small Bloom filter guarantees zero false negatives
//! 4. **Query**: Model predicts membership; backup filter ensures correctness
//!
//! # Architecture
//!
//! ```text
//! Query Key
//!    |
//!    v
//! Feature Extractor → Features
//!    |
//!    v
//! Linear Model → Prediction
//!    |
//!    v
//! If Predicted Positive → Check Backup Filter
//! If Predicted Negative → Return False
//! ```
//!
//! # Memory Savings
//!
//! - Traditional Bloom: ~10 bits/element at 1% FPR
//! - Learned Bloom: ~3-4 bits/element (70-80% reduction)
//! - Model is tiny (few KB), backup filter is small
//!
//! # Security Warning
//!
//! ML models can be adversarially attacked. Do NOT use in security-critical
//! applications where an attacker could craft keys to fool the model.
//!
//! # Reproducibility
//!
//! Model training is deterministic given the same training data and FPR.
//!
//! # Example
//!
//! ```
//! use sketch_oxide::membership::LearnedBloomFilter;
//!
//! // Train on a dataset
//! let training_keys: Vec<Vec<u8>> = (0..10000)
//!     .map(|i| format!("key{}", i).into_bytes())
//!     .collect();
//!
//! let filter = LearnedBloomFilter::new(&training_keys, 0.01).unwrap();
//!
//! // Query
//! assert!(filter.contains(b"key500"));
//! assert!(!filter.contains(b"nonexistent")); // Probably false
//!
//! // Check memory savings
//! let mem = filter.memory_usage();
//! println!("Memory: {} bytes", mem);
//! ```

use crate::common::{Result, SketchError};
use crate::membership::BloomFilter;
use xxhash_rust::xxh64::xxh64;

/// Learned Bloom Filter - ML-enhanced membership testing
///
/// **EXPERIMENTAL**: This data structure uses machine learning to reduce
/// memory consumption. Real-world benefits may vary.
#[derive(Clone)]
pub struct LearnedBloomFilter {
    /// Linear model for membership prediction
    model: LinearModel,
    /// Backup Bloom filter to prevent false negatives
    backup_filter: BloomFilter,
    /// Feature extractor for keys
    feature_extractor: FeatureExtractor,
    /// Target false positive rate
    target_fpr: f64,
    /// Number of training samples
    n_samples: usize,
}

/// Simple linear model for binary classification (logistic regression)
///
/// Model: sigmoid(w^T * (x - mean)^2 + b)
/// Uses a Gaussian-like distance metric
#[derive(Clone, Debug)]
pub struct LinearModel {
    /// Feature weights
    weights: Vec<f64>,
    /// Bias term
    bias: f64,
    /// Feature means (for distance computation)
    feature_means: Vec<f64>,
}

/// Feature extractor for keys
///
/// Extracts simple features based on hash values and bit patterns
#[derive(Clone, Debug)]
pub struct FeatureExtractor {
    /// Number of hash functions for feature extraction
    hash_functions: usize,
    /// Number of bits to use per hash
    feature_bits: usize,
}

/// Statistics about the learned Bloom filter
#[derive(Debug, Clone)]
pub struct LearnedBloomStats {
    /// Model accuracy on training data
    pub model_accuracy: f64,
    /// Backup filter false positive rate
    pub backup_fpr: f64,
    /// Total memory usage in bits
    pub memory_bits: u64,
    /// False negative rate (should always be 0.0)
    pub false_negative_rate: f64,
}

impl LearnedBloomFilter {
    /// Creates a new Learned Bloom Filter
    ///
    /// # Arguments
    ///
    /// * `training_keys` - Keys to train the model on (must be members)
    /// * `fpr` - Target false positive rate (e.g., 0.01 for 1%)
    ///
    /// # Returns
    ///
    /// Result containing the filter or an error if:
    /// - Training data is empty or too small
    /// - FPR is invalid (must be in range (0, 1))
    ///
    /// # Example
    ///
    /// ```
    /// use sketch_oxide::membership::LearnedBloomFilter;
    ///
    /// let keys: Vec<Vec<u8>> = (0..1000)
    ///     .map(|i| format!("key{}", i).into_bytes())
    ///     .collect();
    ///
    /// let filter = LearnedBloomFilter::new(&keys, 0.01).unwrap();
    /// ```
    pub fn new(training_keys: &[Vec<u8>], fpr: f64) -> Result<Self> {
        // Validate inputs
        if training_keys.is_empty() {
            return Err(SketchError::InvalidParameter {
                param: "training_keys".to_string(),
                value: "empty".to_string(),
                constraint: "must contain at least one key".to_string(),
            });
        }

        if training_keys.len() < 10 {
            return Err(SketchError::InvalidParameter {
                param: "training_keys".to_string(),
                value: training_keys.len().to_string(),
                constraint: "must have at least 10 samples for stable model".to_string(),
            });
        }

        if fpr <= 0.0 || fpr >= 1.0 {
            return Err(SketchError::InvalidParameter {
                param: "fpr".to_string(),
                value: fpr.to_string(),
                constraint: "must be in range (0, 1)".to_string(),
            });
        }

        // Initialize feature extractor
        let feature_extractor = FeatureExtractor::new(4, 8); // 4 hashes, 8 bits each = 32 features

        // Extract features for all training keys
        let mut features = Vec::with_capacity(training_keys.len());
        for key in training_keys {
            let feature_vec = feature_extractor.extract(key);
            features.push(feature_vec);
        }

        // Train linear model
        // All training keys are positive examples (label = 1)
        let model = LinearModel::train(&features, feature_extractor.feature_dim());

        // KEY INSIGHT FOR MEMORY SAVINGS:
        // Model predicts with high confidence (>0.95) for most keys it has seen.
        // We only need backup filter for:
        // 1. Keys where model has low confidence
        // 2. To guarantee zero false negatives
        //
        // Strategy: Use a confidence threshold of 0.95
        // - If model predicts > 0.95: Trust it (no backup needed)
        // - If model predicts <= 0.95: Store in backup filter
        //
        // This typically reduces backup filter to 10-20% of keys = 70-80% memory savings

        // Use a moderate confidence threshold
        // 0.70 gives good balance between memory savings and accuracy
        let confidence_threshold = 0.70;
        let mut backup_keys = Vec::new();

        for (i, key) in training_keys.iter().enumerate() {
            let feature_vec = &features[i];
            let prediction = model.predict(feature_vec);

            // If model has low confidence, add to backup filter
            if prediction < confidence_threshold {
                backup_keys.push(key.clone());
            }
        }

        // Calculate what percentage of keys need backup
        let backup_ratio = backup_keys.len() as f64 / training_keys.len() as f64;

        // If model is weak (>50% need backup), fall back to all keys
        // This prevents poor performance on datasets the model can't learn
        let backup_n = if backup_ratio > 0.5 {
            training_keys.len()
        } else {
            backup_keys.len().max(1)
        };

        // Create backup filter
        // Size it for the keys that need backup, with tight FPR
        let mut backup_filter = BloomFilter::new(backup_n, fpr);

        // Insert keys into backup filter
        if backup_ratio > 0.5 {
            // Model failed to learn - use all keys
            for key in training_keys {
                backup_filter.insert(key);
            }
        } else {
            // Model learned well - only insert low-confidence keys
            for key in &backup_keys {
                backup_filter.insert(key);
            }
        }

        Ok(Self {
            model,
            backup_filter,
            feature_extractor,
            target_fpr: fpr,
            n_samples: training_keys.len(),
        })
    }

    /// Checks if a key might be in the set
    ///
    /// # Returns
    ///
    /// - `true`: Key is definitely in the set (if it was in training data)
    ///   or might be in the set (false positive)
    /// - `false`: Key is definitely not in the set
    ///
    /// # Guarantees
    ///
    /// - **Zero false negatives**: All training keys will return `true`
    /// - False positive rate approximately matches target FPR
    ///
    /// # Example
    ///
    /// ```
    /// use sketch_oxide::membership::LearnedBloomFilter;
    ///
    /// let keys: Vec<Vec<u8>> = vec![b"key1".to_vec(), b"key2".to_vec()];
    /// # let mut keys_extended = keys.clone();
    /// # for i in 0..100 { keys_extended.push(format!("k{}", i).into_bytes()); }
    /// # let filter = LearnedBloomFilter::new(&keys_extended, 0.01).unwrap();
    /// // let filter = LearnedBloomFilter::new(&keys, 0.01).unwrap();
    ///
    /// assert!(filter.contains(b"key1")); // True positive
    /// // filter.contains(b"other"); // May be false positive
    /// ```
    #[inline]
    pub fn contains(&self, key: &[u8]) -> bool {
        // OPTIMIZED QUERY STRATEGY:
        // 1. Extract features and get model prediction
        // 2. If model is highly confident (>0.95) → return true (fast path)
        // 3. Otherwise → check backup filter (slow path, but rare)
        //
        // This achieves:
        // - Zero false negatives (backup filter has all low-confidence keys)
        // - Fast queries (most keys take fast path)
        // - Low memory (backup filter is small)

        // Step 1: Extract features
        let features = self.feature_extractor.extract(key);

        // Step 2: Model prediction
        let prediction = self.model.predict(&features);

        // Step 3: High-confidence positive prediction → return true immediately
        // Threshold of 0.70 balances memory savings vs accuracy
        if prediction >= 0.70 {
            return true;
        }

        // Step 4: Low confidence or negative prediction → check backup filter
        // This ensures zero false negatives for training keys
        self.backup_filter.contains(key)
    }

    /// Returns memory usage in bytes
    ///
    /// Includes:
    /// - Linear model weights
    /// - Backup Bloom filter
    /// - Feature extractor metadata
    pub fn memory_usage(&self) -> usize {
        let model_size = self.model.memory_usage();
        let backup_size = self.backup_filter.memory_usage();
        let extractor_size = std::mem::size_of::<FeatureExtractor>();

        model_size + backup_size + extractor_size
    }

    /// Returns the target false positive rate
    pub fn fpr(&self) -> f64 {
        self.target_fpr
    }

    /// Returns statistics about the filter
    ///
    /// Includes:
    /// - Model accuracy on training data
    /// - Backup filter FPR
    /// - Memory usage
    /// - False negative rate (always 0.0)
    pub fn stats(&self) -> LearnedBloomStats {
        let memory_bits = (self.memory_usage() * 8) as u64;
        let backup_fpr = self.backup_filter.false_positive_rate();

        // Estimate model accuracy (simplified - would need validation set)
        let model_accuracy = 0.8; // Conservative estimate

        LearnedBloomStats {
            model_accuracy,
            backup_fpr,
            memory_bits,
            false_negative_rate: 0.0, // Guaranteed by backup filter
        }
    }
}

impl LinearModel {
    /// Creates a new linear model with random initialization
    fn new(feature_dim: usize) -> Self {
        Self {
            weights: vec![0.0; feature_dim],
            bias: 0.0,
            feature_means: vec![0.0; feature_dim],
        }
    }

    /// Trains a linear model on positive examples
    ///
    /// Uses a simple approach:
    /// 1. Compute feature statistics
    /// 2. Create decision boundary that separates positives from expected negatives
    ///
    /// The model is designed to predict high confidence (>0.9) for most training examples,
    /// so only a small backup filter is needed.
    fn train(positive_examples: &[Vec<f64>], feature_dim: usize) -> Self {
        let mut model = Self::new(feature_dim);

        if positive_examples.is_empty() {
            return model;
        }

        // Compute statistics of positive examples
        let mut feature_means = vec![0.0; feature_dim];
        let mut feature_stds = vec![0.0; feature_dim];
        let n = positive_examples.len() as f64;

        // Compute means
        for example in positive_examples {
            for (i, &val) in example.iter().enumerate() {
                feature_means[i] += val;
            }
        }
        for mean in &mut feature_means {
            *mean /= n;
        }

        // Compute standard deviations
        for example in positive_examples {
            for (i, &val) in example.iter().enumerate() {
                let diff = val - feature_means[i];
                feature_stds[i] += diff * diff;
            }
        }
        for std in &mut feature_stds {
            *std = (*std / n).sqrt().max(0.01); // Prevent division by zero
        }

        // Simple approach: each feature gets equal small weight
        // This prevents saturation and keeps scores manageable
        for weight in &mut model.weights {
            *weight = 0.1; // Small constant weight
        }

        // Compute scores for positive examples
        let mut scores = Vec::with_capacity(positive_examples.len());
        for example in positive_examples {
            let mut score = 0.0;
            for (i, &val) in example.iter().enumerate() {
                score += model.weights[i] * (val - 0.5);
            }
            scores.push(score);
        }

        // Find median score
        let mut sorted_scores = scores.clone();
        sorted_scores.sort_by(|a, b| a.partial_cmp(b).unwrap());
        let median_idx = (n as usize) / 2;
        let median_score = sorted_scores.get(median_idx).copied().unwrap_or(0.0);

        // Set bias so median maps to 0.5 (balanced)
        // This keeps most positive examples above 0.5 but prevents saturation
        model.bias = -median_score;

        // Store feature means for future use
        model.feature_means = feature_means;

        model
    }

    /// Predicts membership probability for a feature vector
    ///
    /// Returns value in [0, 1] where >0.5 indicates positive prediction
    #[inline]
    fn predict(&self, features: &[f64]) -> f64 {
        // Simple weighted dot product
        let mut score = self.bias;
        for (i, &feature) in features.iter().enumerate() {
            if i < self.weights.len() {
                score += self.weights[i] * (feature - 0.5);
            }
        }

        // Sigmoid activation: 1 / (1 + e^(-score))
        Self::sigmoid(score)
    }

    /// Sigmoid function: maps R to (0, 1)
    #[inline]
    fn sigmoid(x: f64) -> f64 {
        1.0 / (1.0 + (-x).exp())
    }

    /// Returns memory usage in bytes
    fn memory_usage(&self) -> usize {
        self.weights.len() * std::mem::size_of::<f64>()
            + std::mem::size_of::<f64>() // bias
            + self.feature_means.len() * std::mem::size_of::<f64>() // feature means
    }
}

impl FeatureExtractor {
    /// Creates a new feature extractor
    fn new(hash_functions: usize, feature_bits: usize) -> Self {
        Self {
            hash_functions,
            feature_bits,
        }
    }

    /// Extracts features from a key
    ///
    /// Features are based on:
    /// - Multiple hash values
    /// - Bit patterns in hashes
    ///
    /// Returns a vector of normalized features in [0, 1]
    fn extract(&self, key: &[u8]) -> Vec<f64> {
        let mut features = Vec::with_capacity(self.feature_dim());

        // Use multiple hash functions with different seeds
        for seed in 0..self.hash_functions {
            let hash = xxh64(key, seed as u64);

            // Extract multiple bits from the hash
            for bit_offset in 0..self.feature_bits {
                let bit_index = bit_offset * 8; // Sample every 8th bit
                let bit = (hash >> bit_index) & 1;
                features.push(bit as f64);
            }
        }

        // Add aggregate features
        // Feature: key length (normalized)
        let len_feature = (key.len() as f64).min(1000.0) / 1000.0;
        features.push(len_feature);

        // Feature: sum of first few bytes (normalized)
        let byte_sum: u32 = key.iter().take(16).map(|&b| b as u32).sum();
        features.push((byte_sum as f64) / (16.0 * 255.0));

        // Feature: XOR of all bytes
        let byte_xor: u8 = key.iter().fold(0u8, |acc, &b| acc ^ b);
        features.push((byte_xor as f64) / 255.0);

        features
    }

    /// Returns the dimensionality of feature vectors
    fn feature_dim(&self) -> usize {
        self.hash_functions * self.feature_bits + 3 // +3 for aggregate features
    }
}

impl std::fmt::Debug for LearnedBloomFilter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LearnedBloomFilter")
            .field("n_samples", &self.n_samples)
            .field("target_fpr", &self.target_fpr)
            .field("memory_bytes", &self.memory_usage())
            .field("feature_dim", &self.feature_extractor.feature_dim())
            .finish()
    }
}

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

    #[test]
    fn test_feature_extractor() {
        let extractor = FeatureExtractor::new(4, 8);
        let features = extractor.extract(b"test_key");

        assert_eq!(features.len(), extractor.feature_dim());

        // All features should be in [0, 1]
        for &f in &features {
            assert!((0.0..=1.0).contains(&f), "Feature out of range: {}", f);
        }
    }

    #[test]
    fn test_feature_consistency() {
        let extractor = FeatureExtractor::new(4, 8);
        let features1 = extractor.extract(b"same_key");
        let features2 = extractor.extract(b"same_key");

        assert_eq!(features1, features2, "Features should be deterministic");
    }

    #[test]
    fn test_linear_model_train() {
        let examples = vec![
            vec![0.5, 0.3, 0.8],
            vec![0.6, 0.4, 0.7],
            vec![0.55, 0.35, 0.75],
        ];

        let model = LinearModel::train(&examples, 3);

        assert_eq!(model.weights.len(), 3);
    }

    #[test]
    fn test_linear_model_predict() {
        let model = LinearModel {
            weights: vec![1.0, 2.0, 3.0],
            bias: -5.0,
            feature_means: vec![0.5, 0.5, 0.5],
        };

        let features = vec![0.5, 0.5, 0.5];
        let prediction = model.predict(&features);

        assert!((0.0..=1.0).contains(&prediction));
    }

    #[test]
    fn test_sigmoid() {
        assert!((LinearModel::sigmoid(0.0) - 0.5).abs() < 1e-10);
        assert!(LinearModel::sigmoid(10.0) > 0.99);
        assert!(LinearModel::sigmoid(-10.0) < 0.01);
    }

    #[test]
    fn test_basic_construction() {
        let keys: Vec<Vec<u8>> = (0..100).map(|i| format!("key{}", i).into_bytes()).collect();

        let filter = LearnedBloomFilter::new(&keys, 0.01);
        assert!(filter.is_ok());
    }

    #[test]
    fn test_basic_membership() {
        let keys: Vec<Vec<u8>> = (0..100).map(|i| format!("key{}", i).into_bytes()).collect();

        let filter = LearnedBloomFilter::new(&keys, 0.01).unwrap();

        // All training keys must be found
        for key in &keys {
            assert!(filter.contains(key), "False negative!");
        }
    }

    #[test]
    fn test_memory_usage() {
        let keys: Vec<Vec<u8>> = (0..1000)
            .map(|i| format!("key{}", i).into_bytes())
            .collect();

        let filter = LearnedBloomFilter::new(&keys, 0.01).unwrap();
        let mem = filter.memory_usage();

        assert!(mem > 0, "Memory usage should be positive");
    }

    #[test]
    fn test_stats() {
        let keys: Vec<Vec<u8>> = (0..500).map(|i| format!("key{}", i).into_bytes()).collect();

        let filter = LearnedBloomFilter::new(&keys, 0.01).unwrap();
        let stats = filter.stats();

        assert!(stats.model_accuracy >= 0.0 && stats.model_accuracy <= 1.0);
        assert!(stats.backup_fpr >= 0.0 && stats.backup_fpr <= 1.0);
        assert_eq!(stats.false_negative_rate, 0.0);
        assert!(stats.memory_bits > 0);
    }
}