/*
 * Copyright 2007-2009 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/>.
 */

/**
 * \file sigcore/Hat.h
 */

#ifndef SIGCORE_HAT
#define SIGCORE_HAT

// Sigcore includes
#include "Fraction.h"
#include "Global.h"

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

// Standard includes
#include <cstdlib>

namespace Sigcore
{
/**
 * \class Sigcore::Hat Hat.h sigcore/Hat.h
 * \brief Class to help choose items from a weighted set of items.
 * 
 * Given a set of weighted items to choose from, Hat will pick out a random item from the set.
 */
template<class T> class SIGCORE_EXPORT Hat
{
    public:
        /**
         * Default constructor.
         */
        Hat();
        
        /**
         * Choose an item from the set without removal.
         * 
         * \return The item chosen.
         */
        T pick() const;
        /**
         * Choose an item from the set with removal.
         * 
         * \return The item chosen.
         */
        T take();
        /**
         * Choose an item from the set as well as all identical items.
         * 
         * \return The item chosen.
         */
        T takeAndClear();
        /**
         * Set the weight of an item in the set.
         * 
         * \param key The item to set.
         * \param weight The weight of the item.
         */
        void setCount(const T& key, const int weight);
        /**
         * Add a number of items to the set.
         * 
         * \param key The item to add to the set.
         * \param weight How many of the item to add.
         */
        void add(const T& key, const int weight);
        /**
         * \return The number of distinct items in the set.
         */
        int distinctCount() const;
        /**
         * \return The number of total items in the set.
         */
        int count() const;
        /**
         * \param key The item to get the count of.
         * 
         * \return The amount of \p key in the set.
         */
        int count(const T& key) const;
        /**
         * \param key The item to get the chance of.
         * 
         * \return The chance of choosing \p key.
         */
        Fraction chance(const T& key) const;
    private:
        QMap<T, int> m_items;
        int m_count;
};

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

template<class T> inline T Hat<T>::pick() const
{
    if (!m_count)
        return T();
    int choice = drand48() * m_count;
    QList<T> keys = m_items.keys();
    foreach (const T& key, keys)
    {
        choice -= m_items[key];
        if (choice < 0)
            return key;
    }
    return T();
}

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

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

template<class T> inline void Hat<T>::setCount(const T& key, const int weight)
{
    add(key, weight - m_items[key]);
}

template<class T> inline void Hat<T>::add(const T& key, const int weight)
{
    int nWeight = qMax(weight, -m_items[key]);
    m_items[key] += nWeight;
    m_count += nWeight;
    if (!m_items[key])
        m_items.remove(key);
}

template<class T> inline int Hat<T>::distinctCount() const
{
    return m_items.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_items.contains(key))
        return m_items[key];
    return 0;
}

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

}

#endif