use std::{
    collections::{HashMap, hash_map::Entry},
    hash::Hash,
};

/// Like a set, but keys can be inserted multiple times and, thus,
/// removing a key only removes one occurrence.
pub struct MultiSet<K> {
    inner: HashMap<K, usize>,
}

impl<K> Default for MultiSet<K> {
    fn default() -> Self {
        Self {
            inner: HashMap::new(),
        }
    }
}

impl<K> MultiSet<K>
where
    K: Eq + Hash,
{
    /// Inserts an occurrence of a `key`. If it already exists, an occurrence is
    /// added.
    pub fn insert(&mut self, key: K) {
        match self.inner.entry(key) {
            Entry::Occupied(mut entry) => {
                *entry.get_mut() += 1;
            }
            Entry::Vacant(entry) => {
                entry.insert(1);
            }
        }
    }

    /// Removes an occurrence of a `key`.
    pub fn remove(&mut self, key: K) {
        if let Entry::Occupied(mut entry) = self.inner.entry(key) {
            let &count = entry.get();
            debug_assert_ne!(count, 0);

            if count == 1 {
                entry.remove();
            } else {
                *entry.get_mut() -= 1;
            }
        }
    }

    /// Returns true if `key` is exists in the set.
    pub fn contains(&self, key: &K) -> bool {
        let count = self.inner.get(key);
        debug_assert_ne!(count, Some(&0));

        count.is_some()
    }

    /// Returns an iterator over the *distinct* keys in the set.
    pub fn iter_unique(&self) -> impl Iterator<Item = &K> {
        self.inner.keys()
    }

    /// Returns the length of the *distinct* keys in the set.
    pub fn len_unique(&self) -> usize {
        self.inner.len()
    }

    pub fn is_empty(&self) -> bool {
        self.len_unique() == 0
    }

    pub fn reserve(&mut self, additional: usize) {
        self.inner.reserve(additional);
    }
}

impl<K> Extend<K> for MultiSet<K>
where
    K: Eq + Hash,
{
    fn extend<T: IntoIterator<Item = K>>(&mut self, iter: T) {
        // Copied from std lib implementation
        // See https://docs.rs/hashbrown/0.14.3/src/hashbrown/map.rs.html#6499-6514
        let iter = iter.into_iter();
        let reserve = if self.is_empty() {
            iter.size_hint().0
        } else {
            iter.size_hint().0.div_ceil(2)
        };
        self.reserve(reserve);
        iter.for_each(move |k| {
            self.insert(k);
        });
    }
}

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

    #[test]
    fn insert_and_remove() {
        let mut set = MultiSet::default();

        assert!(!set.contains(&1));
        assert!(!set.contains(&2));
        assert_eq!(set.len_unique(), 0);

        set.insert(1);
        set.insert(1);
        assert!(set.contains(&1));
        assert_eq!(set.len_unique(), 1);

        set.insert(2);
        assert!(set.contains(&1));
        assert!(set.contains(&2));
        assert_eq!(set.len_unique(), 2);

        set.remove(1);
        assert!(set.contains(&1));
        assert!(set.contains(&2));
        assert_eq!(set.len_unique(), 2);

        set.remove(1);
        assert!(!set.contains(&1));
        assert!(set.contains(&2));
        assert_eq!(set.len_unique(), 1);

        set.remove(2);
        assert!(!set.contains(&1));
        assert!(!set.contains(&2));
        assert_eq!(set.len_unique(), 0);
    }
}
