/*
 * Copyright 2007-2008 Ben Boeckel <MathStuf@gmail.com>
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License along
 * with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#ifndef __SIGMOD_HAT__
#define __SIGMOD_HAT__

// Sigmod includes
#include "Global.h"

// Qt includes
#include <QtCore/QtGlobal>
#include <QtCore/QMap>

namespace Sigmod
{
template<class T> class Hat
{
    public:
        Hat();
        
        T pick() const;
        T take();
        T takeAndClear();
        void setCount(const T& key, const int weight);
        void add(const T& key, const int weight);
        int distinctCount() const;
        int count() const;
        int count(const T& key) const;
        double chance(const T& key) const;
        
        int operator[](const T& key) const;
    private:
        QMap<T, int> m_objects;
        int m_count;
};

template<class T> inline Hat<T>::Hat() :
        m_count(0)
{
}

template<class T> inline T Hat<T>::pick() const
{
    int choice = qrand() % m_count;
    foreach (T key, m_objects.keys())
    {
        choice -= m_objects[key];
        if (choice < 0)
            return key;
    }
    return T();
}

template<class T> inline T Hat<T>::take()
{
    T chosen = pick();
    if (!(--m_objects[chosen]))
        m_objects.remove(chosen);
    --m_count;
    return chosen;
}

template<class T> inline T Hat<T>::takeAndClear()
{
    T chosen = pick();
    m_count -= m_objects[chosen];
    m_objects.remove(chosen);
    return chosen;
}

template<class T> inline void Hat<T>::setCount(const T& key, const int weight)
{
    if (m_objects.contains(key))
        m_count -= m_objects[key];
    if (weight)
    {
        m_objects[key] = weight;
        m_count += weight;
    }
    else
        m_objects.remove(key);
}

template<class T> inline void Hat<T>::add(const T& key, const int weight)
{
    Q_ASSERT(0 < weight);
    m_objects[key] += weight;
    m_count += weight;
}

template<class T> inline int Hat<T>::distinctCount() const
{
    return m_objects.size();
}

template<class T> inline int Hat<T>::count() const
{
    return m_count;
}

template<class T> inline int Hat<T>::count(const T& key) const
{
    if (m_objects.contains(key))
        return m_objects[key];
    return 0;
}

template<class T> inline double Hat<T>::chance(const T& key) const
{
    return (double(count(key)) / m_count);
}

template<class T> inline int Hat<T>::operator[](const T& key) const
{
    return count(key);
}

}

#endif