//! # gloomier
//!
//! [Bloom filter](https://en.wikipedia.org/wiki/Bloom_filter) that only computes one hash per operation.

use std::{
    hash::{DefaultHasher, Hash, Hasher},
    marker::PhantomData,
};

use primes::{PrimeSet, Sieve};

/// Used in the unpacking operation to create a bitmask that selects a single
/// bit in a U8. `$idx` must be between 0 and 7, inclusive.
///
/// `$idx = 0` is equal to `0b1000_0000`
/// `$idx = 7` is equal to `0b0000_0001`
macro_rules! mask {
    ($idx:literal) => {
        (1u8 << (7 - $idx))
    };

    ($idx:ident) => {
        (1u8 << (7 - $idx))
    };
}

/// `false_positive_rate` must be in the range (0.0, 1.0) - between 0 and 1, exclusive.
///
/// Helpful when determining what constant to choose, but not at runtime
/// since `K` must be a comptime constant.
pub fn optimal_k(false_positive_rate: f32) -> usize {
    assert!(false_positive_rate > 0.0 && false_positive_rate < 1.0);
    //https://en.wikipedia.org/wiki/Bloom_filter#Optimal_number_of_hash_functions
    let optimal_k = -(false_positive_rate.ln() / std::f32::consts::LN_2);
    optimal_k.ceil() as usize
}

/// Bloom filter implementation that only requires one hash calculation per operation.
/// Generate const parameter `K` with `gloomier::optimal_k` according to your needs.
#[derive(Clone, Debug)]
pub struct Filter<T, const K: usize> {
    arrays: [Vec<bool>; K],
    phantom: PhantomData<T>,
}

// SAFETY: Send and Sync are automatically impled for structs that only contain `[Vec<bool>]`,
// this just makes sure that `PhantomData<T>` can't mess with this
unsafe impl<T, const K: usize> Send for Filter<T, K> {}
unsafe impl<T, const K: usize> Sync for Filter<T, K> {}

impl<T: Hash, const K: usize> Filter<T, K> {
    /// `false_positive_rate`, and `K` should be related according to `gloomier::optimal_k`
    /// for the best performance. Otherwise, the observed false positive rate may not be
    /// in line with your desired threshold.
    pub fn new(expected_entries: usize, false_positive_rate: f32) -> Self {
        let coprimes = Self::coprimes(expected_entries, false_positive_rate);
        Self::new_with_coprimes(coprimes)
    }

    fn new_with_coprimes(coprimes: [usize; K]) -> Self {
        let mut arrays = [const { vec![] }; K];
        for i in 0..K {
            arrays[i] = vec![false; coprimes[i]]
        }
        Self {
            arrays,
            phantom: PhantomData,
        }
    }

    /// insert an element into the Bloom filter. warning, you can't undo this!
    pub fn add(&mut self, entry: T) {
        let indices = self.indices(&entry);
        for (array, index) in self.arrays.iter_mut().zip(indices) {
            array[index] = true
        }
    }

    /// returns false if the element hasn't been inserted,
    /// returns true if the element has been inserted or a collision occurs (with specified false positive rate)
    pub fn check(&self, entry: &T) -> bool {
        self.arrays
            .iter()
            .zip(self.indices(entry))
            .all(|(array, index)| array[index])
    }

    /// generates `K` integers that are pairwise coprime, suitable to use as lengths `k_i`
    fn coprimes(expected_entries: usize, false_positive_rate: f32) -> [usize; K] {
        let bits_per_element = false_positive_rate.ln() * -2.08;
        let total_bits = bits_per_element * expected_entries as f32;
        let min_coprime = (total_bits / K as f32).ceil() as u64;

        let mut pset = Sieve::new();
        let mut primes = [0; K];
        for (slot, prime) in primes
            .iter_mut()
            .zip(pset.iter().skip_while(|&prime| prime < min_coprime))
        {
            // primes are all pairwise coprime :3
            *slot = prime as usize
        }
        primes
    }

    /// for a given `entry`, returns the indices to check or set in the boolean arrays
    fn indices(&self, entry: &T) -> [usize; K] {
        let hash = Self::hash(entry) as usize;
        let iter = self.arrays.iter().map(Vec::len).map(|len| hash % len);
        let mut indices = [0; K];
        for (arr, index) in indices.iter_mut().zip(iter) {
            *arr = index
        }
        indices
    }

    fn hash(entry: &T) -> u64 {
        let mut hasher = DefaultHasher::new();
        entry.hash(&mut hasher);
        hasher.finish()
    }

    /// returns the internal state of the Filter. This includes an array with the
    /// coprimes used in the algorithm, as well as a bitpacked vector with the
    /// actual state.
    ///
    /// see also: [`Filter::with_initial_state()`]
    pub fn into_state(self) -> ([usize; K], Vec<u8>) {
        let mut coprimes = [0; K];
        let mut packed = Vec::with_capacity(self.packed_size());

        for (idx, array) in self.arrays.into_iter().enumerate() {
            coprimes[idx] = array.len();
            for full in array.chunks(8) {
                packed.push(Self::pack_byte(full));
            }
        }

        (coprimes, packed)
    }

    /// create a new Filter with a specific state, retrieved from [`Filter::state()`]
    ///
    /// this method checks the length of the provided state to be sure it matches
    /// what we expect.
    pub fn with_initial_state(coprimes: [usize; K], state: &[u8]) -> Self {
        let mut uninit = Self::new_with_coprimes(coprimes);

        let state_len = state.len();
        let expected_state_len = uninit.packed_size();
        assert_eq!(
            state_len, expected_state_len,
            "expected state of size {state_len} but got state of size {expected_state_len}",
        );

        let mut state_bytes = state.iter();
        for array in uninit.arrays.iter_mut() {
            for full in array.chunks_mut(8) {
                Self::unpack_byte(*state_bytes.next().unwrap(), full);
            }
        }

        uninit
    }

    /// returns the length that `self.arrays` will be when packed into bytes.
    fn packed_size(&self) -> usize {
        self.arrays.iter().map(|arr| arr.len().div_ceil(8)).sum()
    }

    fn unpack_byte(packed: u8, bools: &mut [bool]) {
        assert!(
            bools.len() <= 8,
            "cannot unpack partial byte larger than 8 bits. got {} bools",
            bools.len()
        );

        bools
            .iter_mut()
            .enumerate()
            .for_each(|(idx, bool)| *bool = packed & mask!(idx) > 0);
    }

    fn pack_byte(bools: &[bool]) -> u8 {
        assert!(bools.len() <= 8);

        bools
            .iter()
            .enumerate()
            .fold(0u8, |acc, (idx, flag)| acc | ((*flag as u8) << (7 - idx)))
    }
}

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

    #[test]
    fn it_works() {
        let mut filter = Filter::<_, 4>::new(2, 0.1);
        filter.add(1);
        filter.add(2);
        assert!(filter.check(&1));
        assert!(!filter.check(&3));
    }

    #[test]
    fn correctly_packs_partial_bools() {
        assert_eq!(Filter::<&str, 4>::pack_byte(&[true]), 0b1000_0000);
        assert_eq!(Filter::<&str, 4>::pack_byte(&[false, true]), 0b0100_0000);
        assert_eq!(
            Filter::<&str, 4>::pack_byte(&[false, true, true, true]),
            0b0111_0000
        );

        // All ones/zeros
        assert_eq!(Filter::<&str, 4>::pack_byte(&[true].repeat(8)), 255u8);
        assert_eq!(Filter::<&str, 4>::pack_byte(&[false].repeat(8)), 0u8);
    }

    #[test]
    fn correctly_packs_full_bools() {
        assert_eq!(
            Filter::<&str, 4>::pack_byte(&[true, false, false, false, false, false, false, false]),
            0b1000_0000
        );

        assert_eq!(
            Filter::<&str, 4>::pack_byte(&[false, true, false, false, false, false, false, false]),
            0b0100_0000
        );

        assert_eq!(
            Filter::<&str, 4>::pack_byte(&[false, true, true, true, false, false, false, false]),
            0b0111_0000
        );

        // All ones/zeros
        assert_eq!(Filter::<&str, 4>::pack_byte(&[true].repeat(8)), 255);
        assert_eq!(Filter::<&str, 4>::pack_byte(&[false].repeat(8)), 0);
    }

    #[test]
    fn correctly_packs_unpacks_state() {
        // Create our initial filter state
        let mut filter = Filter::<u8, 4>::new(2, 0.1);
        filter.add(1);
        filter.add(128);
        filter.add(255);

        let (coprimes, state) = filter.into_state();

        // Initialize a Filter with a packed state and make sure it functions
        let revived = Filter::<u8, 4>::with_initial_state(coprimes, &state);
        assert!(revived.check(&1));
        assert!(revived.check(&128));
        assert!(revived.check(&255));
        assert!(!revived.check(&101));

        // Finally, check roundtrip
        let (revived_coprimes, revived_state) = revived.into_state();
        assert_eq!(coprimes, revived_coprimes);
        assert_eq!(state, revived_state)
    }
}