set_theory 1.0.0

//! # MultiSet Implementation
//!
//! A mathematical multiset (bag) implementation allowing duplicate elements.
//!
//! Unlike standard sets, multisets track the multiplicity (count) of each element.

use std::{
    collections::{HashMap, HashSet},
    hash::Hash,
};

use crate::CustomSet;

/// A multiset containing elements with multiplicity.
///
/// `MultiSet` allows duplicate elements, tracking how many times
/// each element appears (its multiplicity).
///
/// # Type Parameters
///
/// * `T` - The type of elements. Must implement `Eq + Hash + Clone`
///
/// # Examples
///
/// ```rust
/// use set_theory::models::MultiSet;
///
/// let ms = MultiSet::from(vec!['a', 'a', 'a', 'b', 'c']);
/// assert_eq!(ms.multiplicity(&'a'), 3);
/// assert_eq!(ms.cardinality(), 5); // Total count including duplicates
/// ```
#[derive(Clone)]
pub struct MultiSet<T: Eq + Hash + Clone> {
    elements: HashMap<T, usize>,
}

impl<T: Eq + Hash + Clone> MultiSet<T> {
    /// Creates an empty multiset.
    ///
    /// # Returns
    ///
    /// A new `MultiSet` with no elements.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let empty = MultiSet::<i32>::empty();
    /// assert!(empty.is_empty());
    /// ```
    pub fn empty() -> Self {
        Self {
            elements: HashMap::new(),
        }
    }

    /// Creates a multiset from an iterable.
    ///
    /// # Arguments
    ///
    /// * `iter` - Any iterable collection of elements
    ///
    /// # Returns
    ///
    /// A new `MultiSet` with element multiplicities tracked.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let ms = MultiSet::new(vec![1, 1, 2, 2, 2, 3]);
    /// assert_eq!(ms.multiplicity(&2), 3);
    /// ```
    pub fn new<I: IntoIterator<Item = T>>(iter: I) -> Self {
        let mut multiset = MultiSet::empty();
        for element in iter {
            multiset.add(element, 1);
        }
        multiset
    }

    /// Returns the multiplicity of an element.
    ///
    /// # Arguments
    ///
    /// * `element` - Reference to the element to check
    ///
    /// # Returns
    ///
    /// The number of times the element appears (0 if not present).
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let ms = MultiSet::from(vec![1, 1, 2, 2, 2, 3]);
    /// assert_eq!(ms.multiplicity(&2), 3);
    /// ```
    pub fn multiplicity(&self, element: &T) -> usize {
        *self.elements.get(element).unwrap_or(&0)
    }

    /// Returns the total cardinality (sum of all multiplicities).
    ///
    /// # Returns
    ///
    /// The total count of all elements including duplicates.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let ms = MultiSet::from(vec![1, 1, 2, 3]);
    /// assert_eq!(ms.cardinality(), 4);
    /// ```
    pub fn cardinality(&self) -> usize {
        self.elements.values().sum()
    }

    /// Returns the number of unique elements.
    ///
    /// # Returns
    ///
    /// The count of distinct elements regardless of multiplicity.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let ms = MultiSet::from(vec![1, 1, 2, 2, 3]);
    /// assert_eq!(ms.unique_count(), 3);
    /// ```
    pub fn unique_count(&self) -> usize {
        self.elements.len()
    }

    /// Checks if the multiset is empty.
    ///
    /// # Returns
    ///
    /// `true` if the multiset contains no elements.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let empty = MultiSet::<i32>::empty();
    /// assert!(empty.is_empty());
    /// ```
    pub fn is_empty(&self) -> bool {
        self.elements.is_empty()
    }

    /// Adds an element with specified multiplicity.
    ///
    /// # Arguments
    ///
    /// * `element` - The element to add
    /// * `count` - Number of times to add (must be > 0)
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let mut ms = MultiSet::<char>::empty();
    /// ms.add('a', 3);
    /// assert_eq!(ms.multiplicity(&'a'), 3);
    /// ```
    pub fn add(&mut self, element: T, count: usize) {
        if count == 0 {
            return;
        }
        *self.elements.entry(element).or_insert(0) += count;
    }

    /// Removes an element with specified multiplicity.
    ///
    /// # Arguments
    ///
    /// * `element` - Reference to the element to remove
    /// * `count` - Number of times to remove
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let mut ms = MultiSet::from(vec!['a', 'a', 'a']);
    /// ms.remove(&'a', 2);
    /// assert_eq!(ms.multiplicity(&'a'), 1);
    /// ```
    pub fn remove(&mut self, element: &T, count: usize) {
        if let Some(current) = self.elements.get_mut(element) {
            if *current <= count {
                self.elements.remove(element);
            } else {
                *current -= count;
            }
        }
    }

    /// Returns the union of this multiset with another.
    ///
    /// # Notation
    ///
    /// P ∪ Q where multiplicity = max(P, Q) for each element
    ///
    /// # Arguments
    ///
    /// * `other` - Reference to the other multiset
    ///
    /// # Returns
    ///
    /// A new `MultiSet` with maximum multiplicities.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let p = MultiSet::from(vec!['a', 'a', 'a', 'c']);
    /// let q = MultiSet::from(vec!['a', 'a', 'b', 'c', 'c']);
    /// let result = p.union(&q);
    /// assert_eq!(result.multiplicity(&'a'), 3);
    /// ```
    pub fn union(&self, other: &Self) -> Self {
        let mut result = MultiSet::empty();
        let all_keys = self
            .elements
            .keys()
            .chain(other.elements.keys())
            .collect::<Vec<_>>()
            .into_iter()
            .collect::<HashSet<_>>();

        for key in all_keys {
            let count =
                (*self.elements.get(key).unwrap_or(&0)).max(*other.elements.get(key).unwrap_or(&0));
            if count > 0 {
                result.add(key.clone(), count);
            }
        }
        result
    }

    /// Returns the intersection of this multiset with another.
    ///
    /// # Notation
    ///
    /// P ∩ Q where multiplicity = min(P, Q) for each element
    ///
    /// # Arguments
    ///
    /// * `other` - Reference to the other multiset
    ///
    /// # Returns
    ///
    /// A new `MultiSet` with minimum multiplicities.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let p = MultiSet::from(vec!['a', 'a', 'a', 'c']);
    /// let q = MultiSet::from(vec!['a', 'a', 'b', 'c', 'c']);
    /// let result = p.intersection(&q);
    /// assert_eq!(result.multiplicity(&'a'), 2);
    /// ```
    pub fn intersection(&self, other: &Self) -> Self {
        let mut result = MultiSet::empty();

        for (key, &count) in &self.elements {
            if let Some(&other_count) = other.elements.get(key) {
                let min_count = count.min(other_count);
                if min_count > 0 {
                    result.add(key.clone(), min_count);
                }
            }
        }
        result
    }

    /// Returns the difference of this multiset with another.
    ///
    /// # Notation
    ///
    /// P - Q where multiplicity = max(0, P - Q) for each element
    ///
    /// # Arguments
    ///
    /// * `other` - Reference to the other multiset
    ///
    /// # Returns
    ///
    /// A new `MultiSet` with subtracted multiplicities (minimum 0).
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let p = MultiSet::from(vec!['a', 'a', 'a', 'b']);
    /// let q = MultiSet::from(vec!['a', 'a', 'b', 'b']);
    /// let result = p.difference(&q);
    /// assert_eq!(result.multiplicity(&'a'), 1);
    /// ```
    pub fn difference(&self, other: &Self) -> Self {
        let mut result = MultiSet::empty();

        for (key, &count) in &self.elements {
            let other_count = other.elements.get(key).unwrap_or(&0);
            let diff = count.saturating_sub(*other_count);
            if diff > 0 {
                result.add(key.clone(), diff);
            }
        }
        result
    }

    /// Returns the sum of this multiset with another.
    ///
    /// # Notation
    ///
    /// P + Q where multiplicity = P + Q for each element
    ///
    /// # Arguments
    ///
    /// * `other` - Reference to the other multiset
    ///
    /// # Returns
    ///
    /// A new `MultiSet` with added multiplicities.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let p = MultiSet::from(vec!['a', 'a', 'b']);
    /// let q = MultiSet::from(vec!['a', 'b', 'b']);
    /// let result = p.sum(&q);
    /// assert_eq!(result.multiplicity(&'a'), 3);
    /// ```
    pub fn sum(&self, other: &Self) -> Self {
        let mut result = MultiSet::empty();
        let all_keys: HashSet<_> = self.elements.keys().chain(other.elements.keys()).collect();

        for key in all_keys {
            let count =
                (*self.elements.get(key).unwrap_or(&0)) + (*other.elements.get(key).unwrap_or(&0));
            if count > 0 {
                result.add(key.clone(), count);
            }
        }
        result
    }

    /// Converts this multiset to a standard set.
    ///
    /// # Returns
    ///
    /// A new `CustomSet` containing only unique elements.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let ms = MultiSet::from(vec![1, 1, 1, 2, 2, 3]);
    /// let set = ms.to_set();
    /// assert_eq!(set.cardinality(), 3);
    /// ```
    pub fn to_set(&self) -> CustomSet<T> {
        CustomSet::new(self.elements.keys().cloned())
    }

    /// Returns an iterator over the unique elements.
    ///
    /// # Returns
    ///
    /// An iterator over references to unique elements.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let ms = MultiSet::from(vec![1, 1, 2, 2, 3]);
    /// let mut count = 0;
    /// for _ in ms.iter() {
    ///     count += 1;
    /// }
    /// assert_eq!(count, 3);
    /// ```
    pub fn iter(&self) -> impl Iterator<Item = &T> {
        self.elements.keys()
    }
}

// Implement From<Vec<T>>
impl<T: Eq + Hash + Clone> From<Vec<T>> for MultiSet<T> {
    /// Creates a MultiSet from a Vec.
    ///
    /// # Arguments
    ///
    /// * `vec` - A vector of elements to create the multiset from
    ///
    /// # Returns
    ///
    /// A new `MultiSet` with element multiplicities tracked.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use set_theory::models::MultiSet;
    ///
    /// let ms = MultiSet::from(vec![1, 1, 2, 2, 2, 3]);
    /// assert_eq!(ms.multiplicity(&2), 3);
    /// ```
    fn from(vec: Vec<T>) -> Self {
        MultiSet::new(vec)
    }
}

// Implement FromIterator
impl<T: Eq + Hash + Clone> FromIterator<T> for MultiSet<T> {
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
        MultiSet::new(iter)
    }
}

// Implement Display
impl<T: Eq + Hash + Clone + std::fmt::Display> std::fmt::Display for MultiSet<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.is_empty() {
            write!(f, "∅")
        } else {
            let entries: Vec<_> = self
                .elements
                .iter()
                .map(|(k, v)| format!("{}:{}", k, v))
                .collect();
            write!(f, "{{{}}}", entries.join(", "))
        }
    }
}

// Implement Debug
impl<T: Eq + Hash + Clone + std::fmt::Debug> std::fmt::Debug for MultiSet<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "MultiSet({:?})", self.elements)
    }
}

// Implement PartialEq
impl<T: Eq + Hash + Clone> PartialEq for MultiSet<T> {
    fn eq(&self, other: &Self) -> bool {
        self.elements == other.elements
    }
}

// Implement Eq
impl<T: Eq + Hash + Clone> Eq for MultiSet<T> {}