// $Id: bag.h,v 1.6 2006-10-15 17:29:09 cactus Exp $ -*-  c++ -*-

#ifndef EAF_BAG_H
#define EAF_BAG_H
#include <iostream>
#include <stdexcept>
#include "list.h"

template<typename T, typename Compare=std::less<T> >
class Bag
{
public:
    typedef T elem_t;
    typedef Compare compare_t;
    
private:
    struct listelem_t
    {
        elem_t elem;
        int    m;

        listelem_t (const elem_t &elem_, int m_ = 1):
            elem (elem_), m (m_) {}
    };
    
    typedef List<listelem_t> list_t;
    typedef typename list_t::iterator list_iterator;
    typedef typename list_t::const_iterator const_list_iterator;
    list_t list;

public:
    class iterator
    {
        list_iterator iter;
        int i;
        
        iterator (const list_iterator &iter_) : iter (iter_), i (iter ? iter->m : 0) {}
        friend class Bag;
        
    public:
        typedef elem_t& reference;
        typedef elem_t* pointer;
        
        reference operator* ()  { return iter->elem; }
        pointer   operator-> () { return iter->elem; }
        
        iterator&      operator++()    { next (); return *this; }
        const iterator operator++(int) { iterator ret = *this; ++*this; return ret; }
        
        bool operator== (const iterator &other) const { return other.iter == iter && other.i == i; }
        bool operator!= (const iterator &other) const { return !(other == *this); }
        
    private:
        void next () {
            if (--i) return;
            
            ++iter;
            i = iter ? iter->m : 0;
        }
    };
    
    class const_iterator
    {
        const_list_iterator iter;
        int i;
        
        const_iterator (const const_list_iterator &iter_) : iter (iter_), i (iter ? iter->m : 0) {}
        friend class Bag;
        
    public:
        const_iterator (const iterator &iter_src): iter (iter_src.iter), i (iter_src.i) {}
        
        typedef elem_t& reference;
        typedef elem_t* pointer;
        
        const reference operator* ()  { return iter->elem; }
        const pointer   operator-> () { return iter->elem; }
        
        const_iterator&      operator++()    { next (); return *this; }
        const const_iterator operator++(int) { const_iterator ret = *this; ++*this; return ret; }
        
        bool operator== (const const_iterator &other) const { return other.iter == iter && other.i == i; }
        bool operator!= (const const_iterator &other) const { return !(other == *this); }
        
    private:
        void next () {
            if (--i) return;
            
            ++iter;
            i = iter ? iter->m : 0;
        }
    };
    
    Bag (){};

    // Add/remove items
    void add (const elem_t &elem);
    void remove (const elem_t &elem);
    void clear () { list.clear (); }

    // Iterate over items (each item is returned according to its multiplicity)
    iterator       begin ()       { return iterator (list.begin ()); }
    const_iterator begin () const { return const_iterator (list.begin ()); }
    
    iterator       end ()       { return iterator (list.end ()); }
    const_iterator end () const { return const_iterator (list.end ()); }

    // Get multiplicity of item
    int operator[] (const elem_t &elem) const;

    // Union
    template<typename T1, typename Compare1>
    friend Bag<T1, Compare1> operator+ (const Bag<T1, Compare1> &lhs,
                                        const Bag<T1, Compare1> &rhs);

    // Copying, assignment
    Bag (const Bag &other) : list (other.list) {}

    Bag& operator= (const Bag &other) {
        if (this == &other) return *this;

        list = other.list;
        return *this;
    }

    // I/O
    template<typename T1, typename Compare1>
    friend std::ostream& operator<< (std::ostream &str, const Bag<T1, Compare1> &bag);

    template<typename T1, typename Compare1>
    friend std::istream& operator>> (std::istream &str, Bag<T1, Compare1> &bag);

private:
    bool find (const elem_t &elem, list_iterator &i) {
        for (i = list.begin ();i != list.end () && compare_t()(i->elem, elem); ++i);
        
        return i != list.end () && i->elem == elem;
    }

    bool find (const elem_t &elem, const_list_iterator &i) const {
        for (i = list.begin ();i != list.end () && compare_t()(i->elem, elem); ++i);
        
        return i != list.end () && i->elem == elem;
    }
};

template<typename T, typename Compare>
int Bag<T, Compare>::operator[] (const elem_t &elem) const
{
    const_list_iterator i;
    if (find (elem, i))
        return i->m;
    else
        return 0;    
}

template<typename T, typename Compare>
void Bag<T, Compare>::add (const elem_t &elem)
{
    list_iterator i;

    if (find (elem, i))
    {
        ++i->m;
    } else {
        list.insert (i, listelem_t (elem, 1));
    }
}

template<typename T, typename Compare>
void Bag<T, Compare>::remove (const elem_t &elem)
{
    list_iterator i;

    if (find (elem, i))
    {
        if (!--i->m)
            list.remove (i);
    } else {
            throw std::runtime_error ("Trying to remove element not in bag");
    }
}

template<typename T, typename Compare>
Bag<T, Compare> operator+ (const Bag<T, Compare> &lhs,
                           const Bag<T, Compare> &rhs)
{
    typedef Bag<T, Compare> bag_t;
    typedef typename bag_t::list_t list_t;
    typedef typename bag_t::listelem_t listelem_t;
    typedef typename bag_t::compare_t compare_t;

    bag_t ret;

    typename list_t::const_iterator i = lhs.list.begin ();
    typename list_t::const_iterator j = rhs.list.begin ();

    while (i != lhs.list.end () || j != rhs.list.end ())
    {
        if (i != lhs.list.end () && j != rhs.list.end () && i->elem == j->elem)
        {
            ret.list.insert (ret.list.end (), listelem_t (i->elem, i->m + j->m));
            ++i;
            ++j;
        } else if (i != lhs.list.end () && (j == rhs.list.end () || compare_t()(i->elem, j->elem))) {
            ret.list.insert (ret.list.end (), listelem_t (i->elem, i->m));
            ++i;
        } else if (j != rhs.list.end () && (i == lhs.list.end () || compare_t()(j->elem, i->elem))) {
            ret.list.insert (ret.list.end (), listelem_t (j->elem, j->m));
            ++j;
        }
    }

    return ret;
}

template<typename T, typename Compare>
std::ostream& operator<< (std::ostream &str, const Bag<T, Compare> &bag)
{
    typedef Bag<T, Compare> bag_t;
    
    str << bag.list.size () << ' ';
    for (typename bag_t::list_t::const_iterator i = bag.list.begin (); i != bag.list.end (); ++i)
        str << i->m << ' ' << i->elem << ' ';
    
    return str;
}

template<typename T, typename Compare>
std::istream& operator>> (std::istream &str, Bag<T, Compare> &bag)
{
    typedef Bag<T, Compare> bag_t;

    bag.clear ();
    
    size_t s;
    for (str >> s; str && s; --s)
    {
        int m;
        typename bag_t::elem_t elem;
        str >> m >> elem;

        if (str)
            bag.list.insert (bag.list.end (), typename bag_t::listelem_t (elem, m));
    }
    
    return str;
}

#endif