range-filters 0.1.0

use rand::Rng;

use crate::Key;
use crate::U64_BITS;
use crate::infix_store::InfixStore;
use crate::utils::longest_common_prefix_length;
use crate::y_fast_trie::YFastTrie;
use std::fmt;

const BASE_IMPLICIT_SIZE: u32 = 10;

/// Diva range filter
///
/// # Arguments
/// * `y_fast_trie` - Y-Fast Trie
/// * `target_size` - Target size
/// * `fpr` - False positive rate
/// * `remainder_size` - Remainder size
///
/// # Example
/// ```rust
/// use range_filters::diva::Diva;
/// let keys = vec![1, 2, 3, 4, 5];
/// let target_size = 1024;
/// let fpr = 0.01;
/// let diva = Diva::new_with_keys(&keys, target_size, fpr);
/// ```
///
/// # Returns
/// * `Diva` - Diva range filter
pub struct Diva {
    y_fast_trie: YFastTrie,
    target_size: usize,
    fpr: f64,
    remainder_size: u8,
}

impl Diva {
    pub fn new(target_size: usize, fpr: f64) -> Self {
        let remainder_size = Self::choose_remainder_size(target_size, fpr);
        const NO_LEVELS: usize = 64;
        Self {
            y_fast_trie: YFastTrie::new(NO_LEVELS),
            target_size,
            fpr,
            remainder_size,
        }
    }

    pub fn new_with_keys(keys: &[Key], target_size: usize, fpr: f64) -> Self {
        let remainder_size = Self::choose_remainder_size(target_size, fpr);
        let mut sorted_keys = keys.to_vec();
        sorted_keys.sort();
        sorted_keys.dedup();

        let mut sampled_keys: Vec<Key> = sorted_keys.iter().step_by(target_size).copied().collect();

        // ensure last key is sampled if not already
        if let Some(&last_key) = sorted_keys.last() {
            if sampled_keys.last() != Some(&last_key) {
                sampled_keys.push(last_key);
            }
        }

        // TODO: make this dynamic based on the key length
        const NO_LEVELS: usize = U64_BITS;
        let mut y_fast_trie = YFastTrie::new_with_keys(&sampled_keys, NO_LEVELS);

        // for each pair of consecutive samples, extract infixes from intermediate keys
        for i in 0..sampled_keys.len().saturating_sub(1) {
            let predecessor = sampled_keys[i];
            let successor = sampled_keys[i + 1];

            // compute extraction parameters from boundary keys
            let (shared_prefix_len, redundant_bits, quotient_bits) =
                Self::get_shared_ignore_implicit_size(&predecessor, &successor, false);

            // find intermediate keys between these samples excluding the samples themselves
            let intermediate_keys: Vec<Key> = sorted_keys
                .iter()
                .filter(|&&k| k > predecessor && k < successor)
                .copied()
                .collect();

            // extract infixes from intermediate keys
            let mut infixes = Vec::new();
            // println!("Processing keys between {} and {} (shared={}, redundant={}, quotient={})",
            //          predecessor, successor, shared_prefix_len, redundant_bits, quotient_bits);
            for key in intermediate_keys {
                let infix = Self::extract_partial_key(
                    key,
                    shared_prefix_len,
                    redundant_bits,
                    quotient_bits,
                    remainder_size,
                );
                // println!("Key {} ({:064b}) -> infix {} ({:064b})",
                //          key, key, infix, infix);
                infixes.push(infix);
            }
            // println!("Infixes before InfixStore creation: {:?}", infixes);

            // create InfixStore and attach to predecessor sample
            if !infixes.is_empty() {
                let infix_store = InfixStore::new_with_infixes(&infixes, remainder_size);
                y_fast_trie.set_infix_store(predecessor, infix_store);
            }
        }

        Self {
            y_fast_trie,
            target_size,
            fpr,
            remainder_size,
        }
    }

    /// insert a new key into DIVA
    pub fn insert(&mut self, key: Key) -> bool {
        let mut rng = rand::thread_rng();
        let chance: usize = rng.gen_range(0..self.target_size);
        let is_sample = chance == 0;

        if !is_sample && self.insert_in_infix(key) {
            return true;
        }
        self.insert_as_sample(key)
    }

    fn insert_in_infix(&mut self, key: Key) -> bool {
        // key should be inserted as a sample if any of the boundary keys are missing
        let (s_low, s_high) = match self
            .y_fast_trie
            .predecessor(key)
            .zip(self.y_fast_trie.successor(key))
        {
            Some(bounds) => bounds,
            None => return false,
        };

        // if the key to be inserted already exists as one of the samples
        if key == s_low || key == s_high {
            return true;
        }

        // extract the infix from the key
        let (shared_bits, redundant_bits, quotient_bits) =
            Self::get_shared_ignore_implicit_size(&s_low, &s_high, false);

        let infix = Self::extract_partial_key(
            key,
            shared_bits,
            redundant_bits,
            quotient_bits,
            self.remainder_size,
        );

        match self.y_fast_trie.get_infix_store(s_low) {
            // attempt to insert in an existing infix store first
            Some(store) => {
                if let Ok(mut store) = store.write() {
                    store.insert(infix)
                } else {
                    false
                }
            }
            // else, create a new store with the infix
            None => {
                let new_store = InfixStore::new_with_infixes(&[infix], self.remainder_size);
                self.y_fast_trie.set_infix_store(s_low, new_store);
                true
            }
        }
    }

    fn insert_as_sample(&mut self, _key: Key) -> bool {
        todo!();
    }

    pub fn delete(&mut self, key: Key) -> bool {
        // key doesn't exist if its out of bounds
        let (s_low, s_high) = match self
            .y_fast_trie
            .predecessor(key)
            .zip(self.y_fast_trie.successor(key))
        {
            Some(bounds) => bounds,
            None => return false,
        };

        // delete from trie if the key is a sample
        if key == s_low || key == s_high {
            return self.delete_sample(key);
        }
        self.delete_from_infix(key, s_low, s_high)
    }

    fn delete_from_infix(&mut self, key: Key, s_low: u64, s_high: u64) -> bool {
        // extract the infix from the key
        let (shared_bits, redundant_bits, quotient_bits) =
            Self::get_shared_ignore_implicit_size(&s_low, &s_high, false);
        let infix = Self::extract_partial_key(
            key,
            shared_bits,
            redundant_bits,
            quotient_bits,
            self.remainder_size,
        );

        // delete from store if the store exists
        match self.y_fast_trie.get_infix_store(s_low) {
            Some(store) => {
                if let Ok(mut store) = store.write() {
                    store.delete(infix)
                } else {
                    false
                }
            }
            None => false,
        }
    }

    fn delete_sample(&mut self, _key: Key) -> bool {
        todo!();
    }

    /// compute redundant bits after first differing bit
    /// redundant bits are consecutive bits with opposite patterns in pred/succ
    /// that can be reconstructed knowing the key is in this range
    fn compute_redundant_bits(key_1: Key, key_2: Key, shared_prefix_len: u8) -> u8 {
        if shared_prefix_len >= 63 {
            return 0;
        }

        let mut redundant_bits = 0u8;

        // start after shared prefix + 1 (skip first differing bit)
        let start_pos = shared_prefix_len + 1;

        for bit_pos in start_pos..64 {
            let shift = 63 - bit_pos;
            let bit_1 = (key_1 >> shift) & 1;
            let bit_2 = (key_2 >> shift) & 1;

            // redundant if pred has 0 and succ has 1 (opposite of first diff bit)
            if bit_1 == 0 && bit_2 == 1 {
                redundant_bits += 1;
            } else {
                break; // stop at first non-redundant bit
            }
        }

        redundant_bits
    }

    /// compute shared prefix, redundant bits, and quotient size
    /// returns: (shared_prefix_len, redundant_bits, quotient_bits)
    pub fn get_shared_ignore_implicit_size(
        key_1: &Key,
        key_2: &Key,
        use_redundant_bits: bool,
    ) -> (u8, u8, u8) {
        // step 1: find shared prefix length (LCP)
        let shared = longest_common_prefix_length(*key_1, *key_2) as u8;

        // step 2: compute redundant bits
        let redundant_bits = if use_redundant_bits {
            Self::compute_redundant_bits(*key_1, *key_2, shared)
        } else {
            0
        };

        // step 3: compute quotient size or aka implicit bits
        let bits_used = shared + redundant_bits;

        if bits_used >= 64 {
            return (shared, redundant_bits, 0);
        }

        let remaining_bits = 64 - bits_used;

        // try to use BASE_IMPLICIT_SIZE quotient bits
        if remaining_bits < BASE_IMPLICIT_SIZE as u8 {
            return (shared, redundant_bits, remaining_bits);
        }

        // extract quotient bits from both keys to check sparsity
        // let shift = remaining_bits - BASE_IMPLICIT_SIZE as u8;
        // let quotient_1 = (key_1 >> shift) & ((1u64 << BASE_IMPLICIT_SIZE) - 1);
        // let quotient_2 = (key_2 >> shift) & ((1u64 << BASE_IMPLICIT_SIZE) - 1);

        // TODO: check if there is a better heuristic for the quotient size
        // add 1 bit if range is sparse (uses < 50% of quotient space)
        // let range_size = quotient_2 - quotient_1 + 1;
        // let quotient_bits = if 2 * range_size < (1u64 << BASE_IMPLICIT_SIZE) {
        //     (BASE_IMPLICIT_SIZE + 1).min(remaining_bits as u32) as u8
        // } else {
        //     BASE_IMPLICIT_SIZE as u8
        // };

        (shared, redundant_bits, BASE_IMPLICIT_SIZE as u8)
    }

    /// extract partial key (infix) from a full key
    /// returns: infix = quotient_bits | remainder_bits

    /// # Arguments
    /// * `key` - The full key to extract from
    /// * `shared_prefix_len` - Number of shared prefix bits to skip
    /// * `redundant_bits` - Number of redundant bits to skip (currently always 0)
    /// * `quotient_bits` - Number of quotient bits to extract (implicit)
    /// * `remainder_bits` - Number of remainder bits to extract (explicit)
    pub fn extract_partial_key(
        key: Key,
        shared_prefix_len: u8,
        redundant_bits: u8,
        quotient_bits: u8,
        remainder_bits: u8,
    ) -> Key {
        let start_bit = shared_prefix_len + redundant_bits;

        if start_bit >= 64 {
            return 0;
        }

        let remaining_bits = 64 - start_bit;
        let bits_to_extract = (quotient_bits + remainder_bits).min(remaining_bits);

        if bits_to_extract == 0 {
            return 0;
        }

        let shift_amount = 64 - start_bit - bits_to_extract;
        let extracted = (key >> shift_amount) & ((1 << bits_to_extract) - 1);

        extracted
    }

    /// get MSB (first differing bit) between predecessor and successor
    pub fn get_msb(key_1: &Key, key_2: &Key) -> u8 {
        let shared = longest_common_prefix_length(*key_1, *key_2);

        if shared >= 64 {
            return 0; // keys are identical
        }

        // extract bit at position 'shared' (first differing bit)
        let bit_pos = 63 - shared;
        ((key_1 >> bit_pos) & 1) as u8
    }

    /// calculate remainder size based on FPR
    /// FPR ≈ 2 / 2^remainder_size
    fn choose_remainder_size(_target_size: usize, fpr: f64) -> u8 {
        // remainder_size = log2(2/FPR) = log2(2) + log2(1/FPR) = 1 - log2(FPR)
        let remainder_size = (1.0 - fpr.log2()).ceil() as u8;
        remainder_size.max(4).min(16) // clamp between 4 and 16 bits
    }

    pub fn pretty_print(&self) {
        print!("{}", self);
    }

    /// Helper function to access InfixStore for a given key with proper error handling
    /// Returns the result of the closure if all required components are found
    fn with_infix_store_for_key<F, R>(&self, key: Key, f: F) -> Option<R>
    where
        F: FnOnce(&InfixStore, Key, Key) -> R,
    {
        let predecessor_infix_store = self.y_fast_trie.predecessor_infix_store(key)?;
        let predecessor_key = self.y_fast_trie.predecessor(key)?;
        let successor_key = self.y_fast_trie.successor(key)?;

        let store = predecessor_infix_store.read().ok()?;
        Some(f(&*store, predecessor_key, successor_key))
    }

    /// Check if any sample exists in the given range [start, end]
    fn has_samples_in_range(&self, start: Key, end: Key) -> bool {
        if let Some(first_sample) = self.y_fast_trie.successor(start) {
            first_sample <= end
        } else {
            false
        }
    }

    /// Point lookup: check if a key exists in the filter
    /// Returns true if key might exist (with FPR), false if definitely doesn't exist
    pub fn contains(&self, key: Key) -> bool {
        if self.y_fast_trie.contains(key) {
            return true;
        }

        // Query InfixStore using helper function
        self.with_infix_store_for_key(key, |store, predecessor_key, successor_key| {
            let result =
                store.point_query(key, predecessor_key, successor_key, self.remainder_size);
            result
        })
        .unwrap_or_else(|| {
            false
        })
    }

    /// Range query: check if any key exists in the given range [start, end] (inclusive)
    /// Returns true if at least one key might exist in the range (with FPR), false if definitely no keys exist
    pub fn range_query(&self, start: Key, end: Key) -> bool {
        if start > end {
            return false;
        }

        // Check if any samples intersect with the range
        if self.has_samples_in_range(start, end) {
            return true;
        }

        // Check all InfixStores
        let mut check_key = start;
        while check_key <= end {
            if let Some(result) =
                self.with_infix_store_for_key(check_key, |store, predecessor_key, successor_key| {
                    // Determine the range to query in this InfixStore
                    let range_start = start.max(predecessor_key + 1);
                    let range_end = end.min(successor_key - 1);

                    if range_start <= range_end {
                        if store.range_query(
                            range_start,
                            range_end,
                            predecessor_key,
                            successor_key,
                            self.remainder_size,
                        ) {
                            return (true, successor_key); // Found match, continue from successor
                        }
                    }

                    (false, successor_key) // No match, but continue from successor
                })
            {
                let (found, successor_key) = result;
                if found {
                    return true;
                }
                check_key = successor_key;
            } else {
                break; // No InfixStore found
            }
        }

        false
    }

    /// Get the number of samples in the Y-Fast Trie
    pub fn sample_count(&self) -> usize {
        self.y_fast_trie.sample_count()
    }
}

impl fmt::Display for Diva {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(
            f,
            "\n╔════════════════════════════════════════════════════════╗"
        )?;
        writeln!(
            f,
            "║                   DIVA RANGE FILTER                    ║"
        )?;
        writeln!(
            f,
            "╚════════════════════════════════════════════════════════╝"
        )?;

        // configuration
        writeln!(f, "\nConfiguration:")?;
        writeln!(f, "  Target size (T):      {}", self.target_size)?;
        writeln!(f, "  False positive rate:  {:.4}%", self.fpr * 100.0)?;
        writeln!(f, "  Remainder size:       {} bits", self.remainder_size)?;

        // stats
        writeln!(f, "\nStatistics:")?;
        writeln!(f, "  Total keys:           {}", self.y_fast_trie.len())?;
        writeln!(
            f,
            "  Sample count:         {}",
            self.y_fast_trie.sample_count()
        )?;
        let avg_bucket_size = if self.y_fast_trie.sample_count() > 0 {
            self.y_fast_trie.len() as f64 / self.y_fast_trie.sample_count() as f64
        } else {
            0.0
        };
        writeln!(f, "  Avg keys per bucket:  {:.1}", avg_bucket_size)?;

        // underlying Y-Fast Trie structure
        write!(f, "{}", self.y_fast_trie)?;

        Ok(())
    }
}

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

    #[test]
    fn test_choose_remainder_size() {
        // FPR = 1% -> remainder_size = 8
        assert_eq!(Diva::choose_remainder_size(1024, 0.01), 8);
        // FPR = 0.1% -> remainder_size = 11
        assert_eq!(Diva::choose_remainder_size(1024, 0.001), 11);
        assert_eq!(Diva::choose_remainder_size(1024, 0.1), 5);
    }

    #[test]
    fn test_get_msb() {
        // first differing bit is 0
        let key1 = 0b0000_0000_0000_0000u64;
        let key2 = 0b1111_1111_1111_1111u64;
        assert_eq!(Diva::get_msb(&key1, &key2), 0);

        // first differing bit is 1
        let key1 = 0b1000_0000_0000_0000u64 << 48;
        let key2 = 0b0111_1111_1111_1111u64 << 48;
        assert_eq!(Diva::get_msb(&key1, &key2), 1);
    }

    #[test]
    fn test_extraction_params() {
        let key1 = 0b0000_0000_1111_0000u64 << 48;
        let key2 = 0b0000_0000_1111_1111u64 << 48;

        let (shared, _redundant, quotient) =
            Diva::get_shared_ignore_implicit_size(&key1, &key2, false);

        assert_eq!(shared, 12); // 12 bits shared prefix
        // assert_eq!(redundant, 0);
        assert!(quotient >= 10); // at least 10 bits for quotient
    }

    #[test]
    fn test_construction_small_dataset() {
        // 100 keys, all fit in one sample
        let keys: Vec<u64> = (0..100).map(|i| i * 1000).collect();
        let diva = Diva::new_with_keys(&keys, 1024, 0.01);

        assert_eq!(diva.target_size, 1024);
        assert_eq!(diva.fpr, 0.01);
        assert_eq!(diva.remainder_size, 8);
    }

    #[test]
    fn test_construction_with_sampling() {
        // 5000 keys - should create ~5 samples
        let keys: Vec<u64> = (0..5000).map(|i| i as u64).collect();
        let target_size = 1024;
        let diva = Diva::new_with_keys(&keys, target_size, 0.01);

        // +1 because we sample the last key too
        let expected_samples = (keys.len() + target_size - 1) / target_size + 1;
        let actual_samples = diva.y_fast_trie.len();

        assert_eq!(actual_samples, expected_samples);
    }

    #[test]
    fn test_construction_single_sample() {
        // 500 keys < target_size -> only 1 sample
        let keys: Vec<u64> = (0..500).map(|i| i * 10).collect();
        let diva = Diva::new_with_keys(&keys, 1024, 0.01);

        assert_eq!(diva.y_fast_trie.sample_count(), 1);
    }

    #[test]
    fn test_point_query_simple() {
        // Simple test case with small numbers
        let keys: Vec<Key> = vec![100, 200, 300, 400, 500, 600, 700];
        let target_size = 1024; // This will create samples at 100, 400, 700
        let diva = Diva::new_with_keys(&keys, target_size, 0.01);

        // println!("Created Diva with {} keys", keys.len());
        // diva.pretty_print();

        // Test that all original keys are found
        for &key in &keys {
            // println!("Testing key: {}", key);
            let found = diva.contains(key);
            assert!(found, "Key {} should be found", key);
        }

        // Test some keys that shouldn't exist
        assert!(!diva.contains(150), "Key 150 should not be found");
        assert!(!diva.contains(250), "Key 250 should not be found");
    }

    #[test]
    fn test_point_query() {
        // Create dataset with keys that will create multiple InfixStores
        let keys: Vec<Key> = (0..3000).map(|i| i * 100).collect();
        let target_size = 1000;
        let diva = Diva::new_with_keys(&keys, target_size, 0.01);

        // Test that all original keys are found
        for &key in &keys {
            assert!(diva.contains(key), "Key {} should be found", key);
        }

        // Test some keys that shouldn't exist
        assert!(!diva.contains(50), "Key 50 should not be found");
        assert!(!diva.contains(150), "Key 150 should not be found");
        assert!(!diva.contains(99999), "Key 99999 should not be found");
    }

    #[test]
    fn test_point_query_with_samples() {
        // Test that sampled keys are found via Y-Fast Trie directly
        let keys: Vec<Key> = vec![100, 200, 300, 400, 500];
        let target_size = 2;
        let diva = Diva::new_with_keys(&keys, target_size, 0.01);

        // Keys 100, 300, 500 should be samples
        // Keys 200, 400 should be in InfixStores
        for &key in &keys {
            assert!(diva.contains(key), "Key {} should be found", key);
        }
    }

    #[test]
    fn test_range_query_simple() {
        // Simple test case with small numbers
        let keys: Vec<Key> = vec![100, 200, 300, 400, 500, 600, 700];
        let target_size = 1024; // This will create one sample with all keys in InfixStore
        let diva = Diva::new_with_keys(&keys, target_size, 0.01);

        // println!("Testing range query on simple dataset");
        // diva.pretty_print();

        // Test range that includes all keys
        let result = diva.range_query(50, 750);
        assert!(result, "Range [50, 750] should include all keys");

        // Test range with no keys
        let result = diva.range_query(10, 50);
        assert!(!result, "Range [10, 50] should include no keys");

        // Test range just below and above existing keys
        let result = diva.range_query(99, 101);
        assert!(result, "Range [99, 101] should include key 100");

        let result = diva.range_query(699, 701);
        assert!(result, "Range [699, 701] should include key 700");

        // Test range that should include existing keys - check specific keys
        // Since we have keys at 200, 300, 400, 500, 600, check if any exist in that range
        let result = diva.range_query(200, 600);
        assert!(result, "Range [200, 600] should include existing keys");
    }

    #[test]
    fn test_range_query_single_infix_store() {
        // Test range query within a single InfixStore
        let keys: Vec<Key> = vec![100, 150, 200, 250, 300, 350, 400];
        let target_size = 2; // Force multiple samples and InfixStores
        let diva = Diva::new_with_keys(&keys, target_size, 0.01);

        // println!("Testing range query within single InfixStore");
        // diva.pretty_print();

        // Test range that spans within one InfixStore
        let result = diva.range_query(180, 280);
        assert!(result, "Range [180, 280] should include keys 200, 250");

        // Test range at boundaries
        let result = diva.range_query(200, 250);
        assert!(result, "Range [200, 250] should include keys 200, 250");

        // Test range with no matching keys
        let result = diva.range_query(125, 149);
        assert!(!result, "Range [125, 149] should include no keys");
    }

    #[test]
    fn test_range_query_multiple_infix_stores() {
        // Test range query spanning multiple InfixStores
        let keys: Vec<Key> = vec![100, 200, 300, 400, 500, 600, 700, 800, 900, 1000];
        let target_size = 2; // Force multiple samples and InfixStores
        let diva = Diva::new_with_keys(&keys, target_size, 0.01);

        // println!("Testing range query across multiple InfixStores");
        // diva.pretty_print();

        // Test range spanning multiple InfixStores
        let result = diva.range_query(250, 750);
        assert!(result, "Range [250, 750] should span multiple InfixStores");

        // Test range that starts before first key and ends after last
        let result = diva.range_query(50, 1050);
        assert!(result, "Range [50, 1050] should include all keys");

        // Test range spanning partial InfixStores
        let result = diva.range_query(150, 850);
        assert!(result, "Range [150, 850] should span partial InfixStores");
    }

    #[test]
    fn test_range_query_with_samples() {
        // Test range queries that include sampled keys
        let keys: Vec<Key> = vec![100, 200, 300, 400, 500];
        let target_size = 2; // This will create samples
        let diva = Diva::new_with_keys(&keys, target_size, 0.01);

        // println!("Testing range query with samples");
        // diva.pretty_print();

        // Test range that includes samples
        let result = diva.range_query(50, 350);
        assert!(
            result,
            "Range [50, 350] should include samples and InfixStore keys"
        );

        // Test range exactly on sample boundaries
        let result = diva.range_query(100, 300);
        assert!(result, "Range [100, 300] should include sample keys");

        // Test range starting from sample
        let result = diva.range_query(300, 450);
        assert!(result, "Range [300, 450] should start from sample");
    }

    #[test]
    fn test_range_query_edge_cases() {
        let keys: Vec<Key> = vec![100, 200, 300, 400, 500];
        let target_size = 1024; // Single InfixStore
        let diva = Diva::new_with_keys(&keys, target_size, 0.01);

        // Test range with start > end (should return false)
        let result = diva.range_query(400, 200);
        assert!(!result, "Invalid range [400, 200] should return false");

        // Test range at exact key boundaries
        let result = diva.range_query(200, 200);
        assert!(result, "Range [200, 200] should include key 200");

        // Test range just missing all keys
        let result = diva.range_query(50, 99);
        assert!(!result, "Range [50, 99] should miss all keys");

        let result = diva.range_query(501, 600);
        assert!(!result, "Range [501, 600] should miss all keys");

        // Test range at boundaries
        let result = diva.range_query(100, 500);
        assert!(result, "Range [100, 500] should include all keys");
    }

    #[test]
    fn test_range_query_large_dataset() {
        // Test with larger dataset to ensure performance
        let keys: Vec<Key> = (0..1000).map(|i| i * 10).collect();
        let target_size = 100; // Create multiple InfixStores
        let diva = Diva::new_with_keys(&keys, target_size, 0.01);

        // Test various range sizes
        let result = diva.range_query(500, 1500);
        assert!(result, "Range [500, 1500] should include many keys");

        let result = diva.range_query(2500, 5000);
        assert!(result, "Range [2500, 5000] should include many keys");

        let result = diva.range_query(8000, 9999);
        assert!(result, "Range [8000, 9999] should include keys near end");

        // Test range with no keys
        let result = diva.range_query(15000, 20000);
        assert!(!result, "Range [15000, 20000] should include no keys");
    }

    #[test]
    fn test_insert_between_samples() {
        let mut diva = Diva::new_with_keys(&[1000, 5000], 1024, 0.01);
        assert!(diva.insert(3000));
    }

    #[test]
    fn test_insert_into_existing_store() {
        let mut diva = Diva::new_with_keys(&[1000, 2000, 5000], 1024, 0.01);
        assert!(diva.insert(2500));
    }

    #[test]
    fn test_insert_duplicate_sample() {
        let mut diva = Diva::new_with_keys(&[1000, 5000], 1024, 0.01);
        assert!(diva.insert(1000));
    }

    #[test]
    fn test_delete_empty_diva() {
        let mut diva = Diva::new(1024, 0.01);
        assert!(!diva.delete(100));
    }

    #[test]
    fn test_delete_non_existent() {
        let mut diva = Diva::new_with_keys(&[1000, 5000], 1024, 0.01);
        assert!(!diva.delete(3000));
    }

    #[test]
    fn test_delete_existing_infix() {
        let mut diva = Diva::new_with_keys(&[1000, 2000, 5000], 1024, 0.01);
        assert!(diva.delete(2000));
        assert!(!diva.delete(2000));
    }

    #[test]
    fn test_insert_delete_roundtrip() {
        let mut diva = Diva::new_with_keys(&[1000, 10000], 1024, 0.01);
        assert!(diva.insert(5000));
        assert!(diva.delete(5000));
        assert!(!diva.delete(5000));
    }

    #[test]
    fn test_multiple_operations() {
        let mut diva = Diva::new_with_keys(&[1000, 10000], 1024, 0.01);
        assert!(diva.insert(2000));
        assert!(diva.insert(3000));
        assert!(diva.delete(2000));
        assert!(!diva.delete(2000));
        assert!(diva.delete(3000));
    }

    #[test]
    fn test_random_operations() {
        use rand::Rng;
        use std::panic;

        let mut diva = Diva::new_with_keys(&[1000, 100000], 1024, 0.01);
        let mut rng = rand::thread_rng();
        let mut inserted = Vec::new();

        for _ in 0..100 {
            let key = rng.gen_range(2000..99000);
            if rng.gen_bool(0.6) {
                let result = panic::catch_unwind(panic::AssertUnwindSafe(|| unsafe {
                    (&mut *(&mut diva as *mut Diva)).insert(key)
                }));
                match result {
                    Ok(true) => inserted.push(key),
                    Ok(false) | Err(_) => {}
                }
            } else if !inserted.is_empty() {
                let idx = rng.gen_range(0..inserted.len());
                let key_to_delete = inserted[idx];
                let delete_result = panic::catch_unwind(panic::AssertUnwindSafe(|| unsafe {
                    (&mut *(&mut diva as *mut Diva)).delete(key_to_delete)
                }));
                if let Ok(true) = delete_result {
                    inserted.swap_remove(idx);
                }
            }
        }
    }
}