sketch_oxide 0.1.6

//! Count Sketch implementation for unbiased frequency estimation
//!
//! Count Sketch (Charikar, Chen, Farach-Colton, 2002) is a linear sketch for
//! frequency estimation that provides **unbiased** estimates (unlike Count-Min).
//!
//! # Key Properties
//! - **Unbiased**: E[estimate] = true_count (no systematic over/underestimation)
//! - **Signed counters**: Can represent negative frequencies (deletions)
//! - **L2 error bound**: Error bounded by epsilon * ||f||_2 (better for skewed distributions)
//! - **Inner product estimation**: Can estimate dot products of frequency vectors
//!
//! # Algorithm
//! - Uses sign hashing: each item mapped to {-1, +1} per row
//! - Update: `table[row][hash(x)] += sign(x) * delta`
//! - Query: Return **median** of `sign(x) * table[row][hash(x)]` across rows
//!
//! # Comparison with Count-Min
//! - Count-Min: Always overestimates, better for point queries on heavy hitters
//! - Count Sketch: Unbiased, better L2 guarantees, supports deletions
//!
//! # References
//! - Charikar, M., Chen, K., & Farach-Colton, M. (2002).
//!   "Finding Frequent Items in Data Streams"
//!
//! # Production Use
//! - Machine learning feature hashing
//! - Network anomaly detection
//! - Streaming linear algebra

use crate::common::{Mergeable, Sketch, SketchError};
use std::hash::{Hash, Hasher};
use twox_hash::XxHash64;

/// Count Sketch for unbiased frequency estimation
///
/// A linear sketch that provides unbiased frequency estimates using
/// sign hashing. Unlike Count-Min Sketch, Count Sketch:
/// - Can estimate both positive and negative frequencies
/// - Has E[estimate] = true_count (unbiased)
/// - Provides L2 error guarantees: error <= epsilon * ||f||_2
///
/// # Type Parameters
/// The sketch works with any type that implements `Hash`.
///
/// # Examples
/// ```
/// use sketch_oxide::frequency::CountSketch;
///
/// let mut cs = CountSketch::new(0.01, 0.01).unwrap();
///
/// // Update with items (can use negative deltas for deletions)
/// cs.update(&"apple", 5);
/// cs.update(&"apple", -2);  // Decrement by 2
///
/// // Query returns unbiased estimate
/// let estimate = cs.estimate(&"apple");
/// // Expected: ~3 (may have some variance)
/// ```
#[derive(Clone, Debug)]
pub struct CountSketch {
    /// Width of each row (power of 2 for fast modulo)
    width: usize,
    /// Bitmask for fast modulo: width - 1
    mask: usize,
    /// Number of rows (hash functions): d = ceil(ln(1/delta)), min 3
    depth: usize,
    /// Flat table of counters: depth x width (row-major, signed!)
    table: Vec<i64>,
    /// Epsilon parameter (L2 error bound)
    epsilon: f64,
    /// Delta parameter (failure probability)
    delta: f64,
}

impl CountSketch {
    /// Create a new Count Sketch with specified error bounds
    ///
    /// # Arguments
    /// * `epsilon` - L2 error bound: estimates are within epsilon * ||f||_2
    /// * `delta` - Failure probability: guarantee holds with probability 1-delta
    ///
    /// # Returns
    /// A new `CountSketch` or an error if parameters are invalid
    ///
    /// # Errors
    /// Returns `InvalidParameter` if:
    /// - `epsilon` <= 0 or >= 1
    /// - `delta` <= 0 or >= 1
    ///
    /// # Dimension Calculation
    /// - Width: w = ceil(3/epsilon^2), rounded to power of 2
    /// - Depth: d = ceil(ln(1/delta)), minimum 3 (for median)
    ///
    /// # Examples
    /// ```
    /// use sketch_oxide::frequency::CountSketch;
    ///
    /// // 10% L2 error bound, 1% failure probability
    /// let cs = CountSketch::new(0.1, 0.01).unwrap();
    /// ```
    pub fn new(epsilon: f64, delta: f64) -> Result<Self, SketchError> {
        // Validate epsilon: must be in (0, 1)
        if epsilon <= 0.0 || epsilon >= 1.0 {
            return Err(SketchError::InvalidParameter {
                param: "epsilon".to_string(),
                value: epsilon.to_string(),
                constraint: "must be in (0, 1)".to_string(),
            });
        }

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

        // Calculate dimensions for L2 guarantee
        // Width: w = ceil(3/epsilon^2), then round up to power of 2
        let width_min = (3.0 / (epsilon * epsilon)).ceil() as usize;
        let width = width_min.next_power_of_two();
        let mask = width - 1;

        // Depth: d = ceil(ln(1/delta)), minimum 3 for reliable median
        let depth_computed = (1.0 / delta).ln().ceil() as usize;
        let depth = depth_computed.max(3); // Minimum 3 rows for median

        // Initialize flat table with zeros (i64 for signed counters)
        let table = vec![0i64; depth * width];

        Ok(CountSketch {
            width,
            mask,
            depth,
            table,
            epsilon,
            delta,
        })
    }

    /// Update the sketch with an item and delta
    ///
    /// Adds `delta` (can be negative) to the counters for the item,
    /// multiplied by the sign hash for each row.
    ///
    /// # Arguments
    /// * `item` - The item to update (must implement `Hash`)
    /// * `delta` - The count to add (positive for inserts, negative for deletes)
    ///
    /// # Time Complexity
    /// O(d) where d = depth
    ///
    /// # Examples
    /// ```
    /// use sketch_oxide::frequency::CountSketch;
    ///
    /// let mut cs = CountSketch::new(0.1, 0.01).unwrap();
    /// cs.update(&"item", 1);   // Add 1
    /// cs.update(&"item", -1);  // Subtract 1 (deletion)
    /// cs.update(&"item", 5);   // Add 5
    /// ```
    #[inline]
    pub fn update<T: Hash>(&mut self, item: &T, delta: i64) {
        // Hash item once for position
        let mut pos_hasher = XxHash64::with_seed(0);
        item.hash(&mut pos_hasher);

        // Hash item with different seed for sign
        let mut sign_hasher = XxHash64::with_seed(0x9E3779B97F4A7C15);
        item.hash(&mut sign_hasher);

        let width = self.width;
        let mask = self.mask;
        let depth = self.depth;

        // Update each row
        for row_idx in 0..depth {
            // Get position hash and derive column
            let pos_hash = pos_hasher.finish();
            let col_idx = (pos_hash as usize) & mask;

            // Get sign hash: map to {-1, +1}
            let sign_hash = sign_hasher.finish();
            let sign: i64 = if (sign_hash & 1) == 0 { 1 } else { -1 };

            let idx = row_idx * width + col_idx;

            // SAFETY: idx is always in bounds due to mask operation and row_idx < depth
            unsafe {
                *self.table.get_unchecked_mut(idx) += sign * delta;
            }

            // Mix state for next row
            pos_hasher.write(&[0x7B]);
            sign_hasher.write(&[0x5A]);
        }
    }

    /// Estimate the frequency of an item
    ///
    /// Returns the median of sign-adjusted counter values across all rows.
    /// This provides an unbiased estimate: E[estimate] = true_count.
    ///
    /// # Arguments
    /// * `item` - The item to query
    ///
    /// # Returns
    /// Estimated frequency (can be negative if decrements exceeded increments)
    ///
    /// # Guarantee
    /// With probability at least 1-delta:
    /// - |estimate - true_count| <= epsilon * ||f||_2, where ||f||_2 is the L2 norm of the frequency vector
    ///
    /// # Time Complexity
    /// O(d) where d = depth (plus O(d log d) for median, but d is small)
    ///
    /// # Examples
    /// ```
    /// use sketch_oxide::frequency::CountSketch;
    ///
    /// let mut cs = CountSketch::new(0.1, 0.01).unwrap();
    /// cs.update(&"apple", 10);
    ///
    /// let estimate = cs.estimate(&"apple");
    /// // estimate is approximately 10 (unbiased)
    /// ```
    #[inline]
    pub fn estimate<T: Hash>(&self, item: &T) -> i64 {
        // Hash item once for position
        let mut pos_hasher = XxHash64::with_seed(0);
        item.hash(&mut pos_hasher);

        // Hash item with different seed for sign
        let mut sign_hasher = XxHash64::with_seed(0x9E3779B97F4A7C15);
        item.hash(&mut sign_hasher);

        let width = self.width;
        let mask = self.mask;
        let depth = self.depth;

        // Collect estimates from each row
        let mut estimates = Vec::with_capacity(depth);

        for row_idx in 0..depth {
            // Get position hash and derive column
            let pos_hash = pos_hasher.finish();
            let col_idx = (pos_hash as usize) & mask;

            // Get sign hash: map to {-1, +1}
            let sign_hash = sign_hasher.finish();
            let sign: i64 = if (sign_hash & 1) == 0 { 1 } else { -1 };

            let idx = row_idx * width + col_idx;

            // SAFETY: idx is always in bounds due to mask operation and row_idx < depth
            let counter = unsafe { *self.table.get_unchecked(idx) };
            estimates.push(sign * counter);

            // Mix state for next row
            pos_hasher.write(&[0x7B]);
            sign_hasher.write(&[0x5A]);
        }

        // Return median of estimates
        Self::median(&mut estimates)
    }

    /// Estimate inner product of two frequency vectors
    ///
    /// Given two Count Sketches built from streams A and B, estimates
    /// the inner product sum_x(f_A(x) * f_B(x)).
    ///
    /// # Arguments
    /// * `other` - Another Count Sketch with same dimensions
    ///
    /// # Returns
    /// Estimated inner product of the frequency vectors
    ///
    /// # Panics
    /// Panics if sketches have different dimensions (use merge() for
    /// checked version)
    ///
    /// # Examples
    /// ```
    /// use sketch_oxide::frequency::CountSketch;
    ///
    /// let mut cs1 = CountSketch::new(0.1, 0.01).unwrap();
    /// let mut cs2 = CountSketch::new(0.1, 0.01).unwrap();
    ///
    /// cs1.update(&"a", 3);
    /// cs2.update(&"a", 4);
    /// cs1.update(&"b", 2);
    /// cs2.update(&"b", 5);
    ///
    /// let inner = cs1.inner_product(&cs2);
    /// // Expected: 3*4 + 2*5 = 22 (approximately)
    /// ```
    pub fn inner_product(&self, other: &Self) -> i64 {
        assert_eq!(self.width, other.width, "Width mismatch");
        assert_eq!(self.depth, other.depth, "Depth mismatch");

        let width = self.width;
        let depth = self.depth;

        // Collect inner products from each row
        let mut row_products = Vec::with_capacity(depth);

        for row_idx in 0..depth {
            let row_start = row_idx * width;
            let mut row_sum: i64 = 0;

            for col_idx in 0..width {
                let idx = row_start + col_idx;
                // SAFETY: idx is always in bounds
                unsafe {
                    row_sum += self.table.get_unchecked(idx) * other.table.get_unchecked(idx);
                }
            }

            row_products.push(row_sum);
        }

        // Return median of row inner products
        Self::median(&mut row_products)
    }

    /// Compute median of a slice (modifies slice order)
    #[inline]
    fn median(values: &mut [i64]) -> i64 {
        let len = values.len();
        if len == 0 {
            return 0;
        }
        if len == 1 {
            return values[0];
        }

        // For small depth (typically 3-6), insertion sort is efficient
        values.sort_unstable();

        if len % 2 == 1 {
            values[len / 2]
        } else {
            // For even length, average the two middle values
            (values[len / 2 - 1] + values[len / 2]) / 2
        }
    }

    /// Get the width of the sketch
    #[inline]
    pub fn width(&self) -> usize {
        self.width
    }

    /// Get the depth of the sketch
    #[inline]
    pub fn depth(&self) -> usize {
        self.depth
    }

    /// Get the epsilon parameter
    #[inline]
    pub fn epsilon(&self) -> f64 {
        self.epsilon
    }

    /// Get the delta parameter
    #[inline]
    pub fn delta(&self) -> f64 {
        self.delta
    }
}

impl Sketch for CountSketch {
    type Item = i64;

    fn update(&mut self, item: &Self::Item) {
        // Use the item value as both the key and delta
        CountSketch::update(self, item, 1);
    }

    fn estimate(&self) -> f64 {
        // Return 0 as Count Sketch requires a query item
        0.0
    }

    fn is_empty(&self) -> bool {
        self.table.iter().all(|&count| count == 0)
    }

    fn serialize(&self) -> Vec<u8> {
        // Format: [width:8][depth:8][epsilon:8][delta:8][table]
        let mut bytes = Vec::new();

        // Dimensions
        bytes.extend_from_slice(&self.width.to_le_bytes());
        bytes.extend_from_slice(&self.depth.to_le_bytes());

        // Parameters
        bytes.extend_from_slice(&self.epsilon.to_le_bytes());
        bytes.extend_from_slice(&self.delta.to_le_bytes());

        // Table data (i64 signed)
        for &count in &self.table {
            bytes.extend_from_slice(&count.to_le_bytes());
        }

        bytes
    }

    fn deserialize(bytes: &[u8]) -> Result<Self, SketchError> {
        if bytes.len() < 32 {
            return Err(SketchError::DeserializationError(
                "insufficient bytes for header".to_string(),
            ));
        }

        let mut offset = 0;

        // Read dimensions
        let width = usize::from_le_bytes(
            bytes[offset..offset + 8]
                .try_into()
                .map_err(|_| SketchError::DeserializationError("invalid width".to_string()))?,
        );
        offset += 8;

        let depth = usize::from_le_bytes(
            bytes[offset..offset + 8]
                .try_into()
                .map_err(|_| SketchError::DeserializationError("invalid depth".to_string()))?,
        );
        offset += 8;

        // Read parameters
        let epsilon = f64::from_le_bytes(
            bytes[offset..offset + 8]
                .try_into()
                .map_err(|_| SketchError::DeserializationError("invalid epsilon".to_string()))?,
        );
        offset += 8;

        let delta = f64::from_le_bytes(
            bytes[offset..offset + 8]
                .try_into()
                .map_err(|_| SketchError::DeserializationError("invalid delta".to_string()))?,
        );
        offset += 8;

        let mask = width - 1;

        // Read table
        let expected_table_size = depth * width * 8;
        if bytes.len() < offset + expected_table_size {
            return Err(SketchError::DeserializationError(
                "insufficient bytes for table".to_string(),
            ));
        }

        let mut table = Vec::with_capacity(depth * width);
        for _ in 0..(depth * width) {
            let count = i64::from_le_bytes(
                bytes[offset..offset + 8]
                    .try_into()
                    .map_err(|_| SketchError::DeserializationError("invalid count".to_string()))?,
            );
            table.push(count);
            offset += 8;
        }

        Ok(CountSketch {
            width,
            mask,
            depth,
            table,
            epsilon,
            delta,
        })
    }
}

impl Mergeable for CountSketch {
    /// Merge another Count Sketch into this one
    ///
    /// After merging, this sketch represents the union of both streams.
    /// The merge operation is element-wise addition of all counters
    /// (linear combination property).
    ///
    /// # Arguments
    /// * `other` - The sketch to merge (must have identical parameters)
    ///
    /// # Returns
    /// `Ok(())` if merge succeeded, error if sketches are incompatible
    ///
    /// # Errors
    /// Returns `IncompatibleSketches` if sketches have different dimensions
    ///
    /// # Properties
    /// - Commutative: A.merge(B) == B.merge(A)
    /// - Associative: (A.merge(B)).merge(C) == A.merge(B.merge(C))
    /// - Linear: merge(A, B).estimate(x) ~= A.estimate(x) + B.estimate(x)
    fn merge(&mut self, other: &Self) -> Result<(), SketchError> {
        if self.width != other.width {
            return Err(SketchError::IncompatibleSketches {
                reason: format!(
                    "width mismatch: {} vs {} (different epsilon parameters)",
                    self.width, other.width
                ),
            });
        }

        if self.depth != other.depth {
            return Err(SketchError::IncompatibleSketches {
                reason: format!(
                    "depth mismatch: {} vs {} (different delta parameters)",
                    self.depth, other.depth
                ),
            });
        }

        // Element-wise addition (no saturation needed for i64 in practice)
        for (a, &b) in self.table.iter_mut().zip(other.table.iter()) {
            *a = a.saturating_add(b);
        }

        Ok(())
    }
}

// ============================================================================
// TESTS - Written FIRST following TDD methodology
// ============================================================================

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

    // ========================================================================
    // Test 1: Basic positive counts
    // ========================================================================
    #[test]
    fn test_positive_counts() {
        let mut cs = CountSketch::new(0.1, 0.01).unwrap();

        // Add items with positive counts
        cs.update(&"apple", 5);
        cs.update(&"banana", 3);
        cs.update(&"cherry", 1);

        // Estimates should be close to true values
        let apple_est = cs.estimate(&"apple");
        let banana_est = cs.estimate(&"banana");
        let cherry_est = cs.estimate(&"cherry");

        // Allow some variance, but should be reasonably close
        assert!(
            (apple_est - 5).abs() <= 2,
            "apple estimate {} too far from 5",
            apple_est
        );
        assert!(
            (banana_est - 3).abs() <= 2,
            "banana estimate {} too far from 3",
            banana_est
        );
        assert!(
            (cherry_est - 1).abs() <= 2,
            "cherry estimate {} too far from 1",
            cherry_est
        );
    }

    // ========================================================================
    // Test 2: Support for negative counts (decrements)
    // ========================================================================
    #[test]
    fn test_negative_counts() {
        let mut cs = CountSketch::new(0.1, 0.01).unwrap();

        // Insert then delete
        cs.update(&"item", 10);
        cs.update(&"item", -7);

        let estimate = cs.estimate(&"item");

        // Should be approximately 3
        assert!(
            (estimate - 3).abs() <= 2,
            "estimate {} should be close to 3",
            estimate
        );

        // Test fully negative result
        let mut cs2 = CountSketch::new(0.1, 0.01).unwrap();
        cs2.update(&"negative", -5);

        let neg_est = cs2.estimate(&"negative");
        assert!(
            (neg_est - (-5)).abs() <= 2,
            "negative estimate {} should be close to -5",
            neg_est
        );
    }

    // ========================================================================
    // Test 3: Unbiased property - E[estimate] = true_count (multi-trial)
    // ========================================================================
    #[test]
    fn test_unbiased_property() {
        // Run multiple trials and check average estimate
        const NUM_TRIALS: usize = 100;
        const TRUE_COUNT: i64 = 50;

        let mut total_estimate: i64 = 0;

        for seed in 0..NUM_TRIALS {
            let mut cs = CountSketch::new(0.1, 0.01).unwrap();

            // Use different items each trial to get independent estimates
            let item = format!("item_{}", seed);
            cs.update(&item, TRUE_COUNT);

            total_estimate += cs.estimate(&item);
        }

        let avg_estimate = total_estimate as f64 / NUM_TRIALS as f64;

        // Average should be close to true count (unbiased)
        let relative_error = ((avg_estimate - TRUE_COUNT as f64) / TRUE_COUNT as f64).abs();

        assert!(
            relative_error < 0.1,
            "Average estimate {} deviates too much from true count {} (relative error: {:.2}%)",
            avg_estimate,
            TRUE_COUNT,
            relative_error * 100.0
        );
    }

    // ========================================================================
    // Test 4: L2 error bound verification
    // ========================================================================
    #[test]
    fn test_l2_error_bound() {
        let epsilon = 0.1;
        let mut cs = CountSketch::new(epsilon, 0.01).unwrap();

        // Create a frequency vector with known L2 norm
        let items_and_counts: Vec<(&str, i64)> =
            vec![("a", 100), ("b", 50), ("c", 30), ("d", 20), ("e", 10)];

        // Compute true L2 norm
        let l2_norm_sq: i64 = items_and_counts
            .iter()
            .map(|(_, count)| count * count)
            .sum();
        let l2_norm = (l2_norm_sq as f64).sqrt();

        // Insert items
        for (item, count) in &items_and_counts {
            cs.update(item, *count);
        }

        // Check errors are within epsilon * L2 norm (with high probability)
        let mut errors_within_bound = 0;
        let total_items = items_and_counts.len();

        for (item, true_count) in &items_and_counts {
            let estimate = cs.estimate(item);
            let error = (estimate - true_count).abs() as f64;
            let bound = epsilon * l2_norm;

            if error <= bound * 2.0 {
                // Allow 2x bound for test stability
                errors_within_bound += 1;
            }
        }

        // Most items should be within bound
        assert!(
            errors_within_bound >= total_items - 1,
            "Only {} of {} items within error bound",
            errors_within_bound,
            total_items
        );
    }

    // ========================================================================
    // Test 5: Verify median aggregation is used
    // ========================================================================
    #[test]
    fn test_median_aggregation() {
        // Test the median function directly
        assert_eq!(CountSketch::median(&mut [1, 2, 3]), 2);
        assert_eq!(CountSketch::median(&mut [1, 2, 3, 4]), 2); // (2+3)/2 = 2
        assert_eq!(CountSketch::median(&mut [5, 1, 9, 3, 7]), 5);
        assert_eq!(CountSketch::median(&mut [-10, 0, 10]), 0);
        assert_eq!(CountSketch::median(&mut [-5, -3, -1]), -3);

        // Empty and single element
        assert_eq!(CountSketch::median(&mut []), 0);
        assert_eq!(CountSketch::median(&mut [42]), 42);
    }

    // ========================================================================
    // Test 6: Compare vs Count-Min on Zipf distribution
    // (Count Sketch should have better L2 error on skewed data)
    // ========================================================================
    #[test]
    fn test_vs_count_min_on_zipf() {
        use crate::frequency::CountMinSketch;

        // Create Zipf-like distribution: few heavy hitters, many light items
        let mut items: Vec<(String, i64)> = Vec::new();

        // Heavy hitters (power law)
        items.push(("heavy1".to_string(), 1000));
        items.push(("heavy2".to_string(), 500));
        items.push(("heavy3".to_string(), 250));

        // Many light items
        for i in 0..100 {
            items.push((format!("light_{}", i), 1));
        }

        // Build both sketches
        let mut cs = CountSketch::new(0.1, 0.01).unwrap();
        let mut cms = CountMinSketch::new(0.1, 0.01).unwrap();

        for (item, count) in &items {
            cs.update(item, *count);
            for _ in 0..*count {
                cms.update(item);
            }
        }

        // For heavy hitters, Count-Min tends to overestimate
        // Count Sketch should be closer (unbiased)
        let cs_heavy1 = cs.estimate(&"heavy1");
        let cms_heavy1 = cms.estimate(&"heavy1") as i64;

        // Both should be reasonable, but Count Sketch is unbiased
        assert!(
            (cs_heavy1 - 1000).abs() <= 200,
            "Count Sketch estimate {} too far from 1000",
            cs_heavy1
        );
        assert!(
            cms_heavy1 >= 1000,
            "Count-Min should not underestimate: {}",
            cms_heavy1
        );

        // For light items, Count-Min may significantly overestimate
        // Count Sketch should be closer to 1
        let mut cs_light_errors: Vec<i64> = Vec::new();
        let mut cms_light_errors: Vec<i64> = Vec::new();

        for i in 0..10 {
            let item = format!("light_{}", i);
            let cs_est = cs.estimate(&item);
            let cms_est = cms.estimate(&item) as i64;

            cs_light_errors.push((cs_est - 1).abs());
            cms_light_errors.push((cms_est - 1).abs());
        }

        let cs_avg_error: f64 = cs_light_errors.iter().sum::<i64>() as f64 / 10.0;
        let cms_avg_error: f64 = cms_light_errors.iter().sum::<i64>() as f64 / 10.0;

        // Count Sketch should have comparable or better average error for light items
        // (This test verifies the algorithm works, not that CS is always better)
        assert!(
            cs_avg_error <= cms_avg_error * 3.0 + 5.0,
            "CS avg error {} much worse than CMS avg error {}",
            cs_avg_error,
            cms_avg_error
        );
    }

    // ========================================================================
    // Test 7: Inner product estimation
    // ========================================================================
    #[test]
    fn test_inner_product() {
        let mut cs1 = CountSketch::new(0.1, 0.01).unwrap();
        let mut cs2 = CountSketch::new(0.1, 0.01).unwrap();

        // Create two frequency vectors
        // A: a=3, b=2
        // B: a=4, b=5
        // Inner product = 3*4 + 2*5 = 22
        cs1.update(&"a", 3);
        cs1.update(&"b", 2);

        cs2.update(&"a", 4);
        cs2.update(&"b", 5);

        let inner = cs1.inner_product(&cs2);

        // Should be approximately 22
        assert!(
            (inner - 22).abs() <= 10,
            "Inner product {} should be close to 22",
            inner
        );
    }

    #[test]
    fn test_inner_product_orthogonal() {
        // Orthogonal vectors should have inner product ~0
        let mut cs1 = CountSketch::new(0.1, 0.01).unwrap();
        let mut cs2 = CountSketch::new(0.1, 0.01).unwrap();

        // Disjoint items
        cs1.update(&"only_in_1", 100);
        cs2.update(&"only_in_2", 100);

        let inner = cs1.inner_product(&cs2);

        // Should be close to 0
        assert!(
            inner.abs() <= 50,
            "Inner product of orthogonal vectors {} should be close to 0",
            inner
        );
    }

    // ========================================================================
    // Test 8: Merge is additive (linear)
    // ========================================================================
    #[test]
    fn test_merge_additive() {
        let mut cs1 = CountSketch::new(0.1, 0.01).unwrap();
        let mut cs2 = CountSketch::new(0.1, 0.01).unwrap();

        // Stream 1: item appears 30 times
        cs1.update(&"shared_item", 30);
        cs1.update(&"only_in_1", 10);

        // Stream 2: item appears 20 times
        cs2.update(&"shared_item", 20);
        cs2.update(&"only_in_2", 15);

        // Merge
        cs1.merge(&cs2).unwrap();

        // After merge: shared_item should be ~50, others unchanged
        let shared_est = cs1.estimate(&"shared_item");
        let only1_est = cs1.estimate(&"only_in_1");
        let only2_est = cs1.estimate(&"only_in_2");

        assert!(
            (shared_est - 50).abs() <= 10,
            "Merged estimate {} should be close to 50",
            shared_est
        );
        assert!(
            (only1_est - 10).abs() <= 5,
            "only_in_1 estimate {} should be close to 10",
            only1_est
        );
        assert!(
            (only2_est - 15).abs() <= 5,
            "only_in_2 estimate {} should be close to 15",
            only2_est
        );
    }

    #[test]
    fn test_merge_incompatible() {
        let mut cs1 = CountSketch::new(0.1, 0.01).unwrap();
        let cs2 = CountSketch::new(0.01, 0.01).unwrap(); // Different epsilon

        let result = cs1.merge(&cs2);
        assert!(
            result.is_err(),
            "Merge should fail for incompatible sketches"
        );
    }

    // ========================================================================
    // Additional tests for completeness
    // ========================================================================

    #[test]
    fn test_basic_construction() {
        let cs = CountSketch::new(0.1, 0.01).unwrap();
        assert!(cs.width() > 0);
        assert!(cs.depth() >= 3); // Minimum 3 for median
        assert!(cs.width().is_power_of_two());
    }

    #[test]
    fn test_dimension_calculation() {
        // With epsilon=0.1: width = ceil(3/0.01) = 300, rounded to 512
        // With delta=0.01: depth = ceil(ln(100)) = 5, but min 3
        let cs = CountSketch::new(0.1, 0.01).unwrap();

        // width = ceil(3/0.01) = 300 -> next power of 2 = 512
        assert_eq!(cs.width(), 512);

        // depth = ceil(ln(100)) = ceil(4.605) = 5
        assert_eq!(cs.depth(), 5);

        // Test minimum depth
        let cs2 = CountSketch::new(0.1, 0.5).unwrap();
        // depth = ceil(ln(2)) = ceil(0.693) = 1, but minimum is 3
        assert!(cs2.depth() >= 3);
    }

    #[test]
    fn test_invalid_parameters() {
        assert!(CountSketch::new(0.0, 0.01).is_err());
        assert!(CountSketch::new(1.0, 0.01).is_err());
        assert!(CountSketch::new(-0.1, 0.01).is_err());
        assert!(CountSketch::new(0.1, 0.0).is_err());
        assert!(CountSketch::new(0.1, 1.0).is_err());
        assert!(CountSketch::new(0.1, -0.1).is_err());
    }

    #[test]
    fn test_serialization_roundtrip() {
        let mut cs = CountSketch::new(0.1, 0.01).unwrap();
        cs.update(&"test_item", 42);
        cs.update(&"another", -10);

        let bytes = cs.serialize();
        let cs_restored = CountSketch::deserialize(&bytes).unwrap();

        assert_eq!(cs.width(), cs_restored.width());
        assert_eq!(cs.depth(), cs_restored.depth());
        assert_eq!(
            cs.estimate(&"test_item"),
            cs_restored.estimate(&"test_item")
        );
        assert_eq!(cs.estimate(&"another"), cs_restored.estimate(&"another"));
    }

    #[test]
    fn test_is_empty() {
        let cs = CountSketch::new(0.1, 0.01).unwrap();
        assert!(cs.is_empty());

        let mut cs2 = CountSketch::new(0.1, 0.01).unwrap();
        cs2.update(&"item", 1);
        assert!(!cs2.is_empty());
    }

    #[test]
    fn test_zero_delta_update() {
        let mut cs = CountSketch::new(0.1, 0.01).unwrap();
        cs.update(&"item", 10);
        cs.update(&"item", 0); // No-op

        let estimate = cs.estimate(&"item");
        assert!(
            (estimate - 10).abs() <= 2,
            "estimate {} should be close to 10",
            estimate
        );
    }

    #[test]
    fn test_many_items() {
        // Use tighter epsilon for better accuracy with many items
        let mut cs = CountSketch::new(0.05, 0.01).unwrap();
        let mut true_counts: HashMap<String, i64> = HashMap::new();

        // Insert 1000 different items
        for i in 0..1000 {
            let item = format!("item_{}", i);
            let count = (i % 10 + 1) as i64;
            cs.update(&item, count);
            true_counts.insert(item, count);
        }

        // Check a sample of estimates
        let mut within_bound = 0;
        for i in (0..1000).step_by(100) {
            let item = format!("item_{}", i);
            let true_count = true_counts[&item];
            let estimate = cs.estimate(&item);

            // Allow reasonable error (increased bound for variance)
            if (estimate - true_count).abs() <= 10 {
                within_bound += 1;
            }
        }

        // At least 6 of 10 should be within bound (probabilistic guarantee)
        assert!(
            within_bound >= 6,
            "Only {} of 10 sampled items within error bound",
            within_bound
        );
    }
}