/* boost random/discrete_distribution.hpp header file
 *
 * Copyright Steven Watanabe 2009-2011
 * Distributed under the Boost Software License, Version 1.0. (See
 * accompanying file LICENSE_1_0.txt or copy at
 * http://www.boost.org/LICENSE_1_0.txt)
 *
 * See http://www.boost.org for most recent version including documentation.
 *
 * $Id$
 */

#ifndef BOOST_RANDOM_DISCRETE_DISTRIBUTION_HPP_INCLUDED
#define BOOST_RANDOM_DISCRETE_DISTRIBUTION_HPP_INCLUDED

#include <vector>
#include <limits>
#include <numeric>
#include <utility>
#include <iterator>
#include <boost/assert.hpp>
#include <boost/random/uniform_01.hpp>
#include <boost/random/uniform_int_distribution.hpp>
#include <boost/random/detail/config.hpp>
#include <boost/random/detail/operators.hpp>
#include <boost/random/detail/vector_io.hpp>

#ifndef BOOST_NO_CXX11_HDR_INITIALIZER_LIST
#include <initializer_list>
#endif

#include <boost/range/begin.hpp>
#include <boost/range/end.hpp>

#include <boost/random/detail/disable_warnings.hpp>

namespace boost {
namespace random {
namespace detail {

template<class IntType, class WeightType>
struct integer_alias_table {
    WeightType get_weight(IntType bin) const {
        WeightType result = _average;
        if(bin < _excess) ++result;
        return result;
    }
    template<class Iter>
    WeightType init_average(Iter begin, Iter end) {
        WeightType weight_average = 0;
        IntType excess = 0;
        IntType n = 0;
        // weight_average * n + excess == current partial sum
        // This is a bit messy, but it's guaranteed not to overflow
        for(Iter iter = begin; iter != end; ++iter) {
            ++n;
            if(*iter < weight_average) {
                WeightType diff = weight_average - *iter;
                weight_average -= diff / n;
                if(diff % n > excess) {
                    --weight_average;
                    excess += n - diff % n;
                } else {
                    excess -= diff % n;
                }
            } else {
                WeightType diff = *iter - weight_average;
                weight_average += diff / n;
                if(diff % n < n - excess) {
                    excess += diff % n;
                } else {
                    ++weight_average;
                    excess -= n - diff % n;
                }
            }
        }
        _alias_table.resize(static_cast<std::size_t>(n));
        _average = weight_average;
        _excess = excess;
        return weight_average;
    }
    void init_empty()
    {
        _alias_table.clear();
        _alias_table.push_back(std::make_pair(static_cast<WeightType>(1),
                                              static_cast<IntType>(0)));
        _average = static_cast<WeightType>(1);
        _excess = static_cast<IntType>(0);
    }
    bool operator==(const integer_alias_table& other) const
    {
        return _alias_table == other._alias_table &&
            _average == other._average && _excess == other._excess;
    }
    static WeightType normalize(WeightType val, WeightType /* average */)
    {
        return val;
    }
    static void normalize(std::vector<WeightType>&) {}
    template<class URNG>
    WeightType test(URNG &urng) const
    {
        return uniform_int_distribution<WeightType>(0, _average)(urng);
    }
    bool accept(IntType result, WeightType val) const
    {
        return result < _excess || val < _average;
    }
    static WeightType try_get_sum(const std::vector<WeightType>& weights)
    {
        WeightType result = static_cast<WeightType>(0);
        for(typename std::vector<WeightType>::const_iterator
                iter = weights.begin(), end = weights.end();
            iter != end; ++iter)
        {
            if((std::numeric_limits<WeightType>::max)() - result > *iter) {
                return static_cast<WeightType>(0);
            }
            result += *iter;
        }
        return result;
    }
    template<class URNG>
    static WeightType generate_in_range(URNG &urng, WeightType max)
    {
        return uniform_int_distribution<WeightType>(
            static_cast<WeightType>(0), max-1)(urng);
    }
    typedef std::vector<std::pair<WeightType, IntType> > alias_table_t;
    alias_table_t _alias_table;
    WeightType _average;
    IntType _excess;
};

template<class IntType, class WeightType>
struct real_alias_table {
    WeightType get_weight(IntType) const
    {
        return WeightType(1.0);
    }
    template<class Iter>
    WeightType init_average(Iter first, Iter last)
    {
        std::size_t size = std::distance(first, last);
        WeightType weight_sum =
            std::accumulate(first, last, static_cast<WeightType>(0));
        _alias_table.resize(size);
        return weight_sum / size;
    }
    void init_empty()
    {
        _alias_table.clear();
        _alias_table.push_back(std::make_pair(static_cast<WeightType>(1),
                                              static_cast<IntType>(0)));
    }
    bool operator==(const real_alias_table& other) const
    {
        return _alias_table == other._alias_table;
    }
    static WeightType normalize(WeightType val, WeightType average)
    {
        return val / average;
    }
    static void normalize(std::vector<WeightType>& weights)
    {
        WeightType sum =
            std::accumulate(weights.begin(), weights.end(),
                            static_cast<WeightType>(0));
        for(typename std::vector<WeightType>::iterator
                iter = weights.begin(),
                end = weights.end();
            iter != end; ++iter)
        {
            *iter /= sum;
        }
    }
    template<class URNG>
    WeightType test(URNG &urng) const
    {
        return uniform_01<WeightType>()(urng);
    }
    bool accept(IntType, WeightType) const
    {
        return true;
    }
    static WeightType try_get_sum(const std::vector<WeightType>& /* weights */)
    {
        return static_cast<WeightType>(1);
    }
    template<class URNG>
    static WeightType generate_in_range(URNG &urng, WeightType)
    {
        return uniform_01<WeightType>()(urng);
    }
    typedef std::vector<std::pair<WeightType, IntType> > alias_table_t;
    alias_table_t _alias_table;
};

template<bool IsIntegral>
struct select_alias_table;

template<>
struct select_alias_table<true> {
    template<class IntType, class WeightType>
    struct apply {
        typedef integer_alias_table<IntType, WeightType> type;
    };
};

template<>
struct select_alias_table<false> {
    template<class IntType, class WeightType>
    struct apply {
        typedef real_alias_table<IntType, WeightType> type;
    };
};

}

/**
 * The class @c discrete_distribution models a \random_distribution.
 * It produces integers in the range [0, n) with the probability
 * of producing each value is specified by the parameters of the
 * distribution.
 */
template<class IntType = int, class WeightType = double>
class discrete_distribution {
public:
    typedef WeightType input_type;
    typedef IntType result_type;

    class param_type {
    public:

        typedef discrete_distribution distribution_type;

        /**
         * Constructs a @c param_type object, representing a distribution
         * with \f$p(0) = 1\f$ and \f$p(k|k>0) = 0\f$.
         */
        param_type() : _probabilities(1, static_cast<WeightType>(1)) {}
        /**
         * If @c first == @c last, equivalent to the default constructor.
         * Otherwise, the values of the range represent weights for the
         * possible values of the distribution.
         */
        template<class Iter>
        param_type(Iter first, Iter last) : _probabilities(first, last)
        {
            normalize();
        }
#ifndef BOOST_NO_CXX11_HDR_INITIALIZER_LIST
        /**
         * If wl.size() == 0, equivalent to the default constructor.
         * Otherwise, the values of the @c initializer_list represent
         * weights for the possible values of the distribution.
         */
        param_type(const std::initializer_list<WeightType>& wl)
          : _probabilities(wl)
        {
            normalize();
        }
#endif
        /**
         * If the range is empty, equivalent to the default constructor.
         * Otherwise, the elements of the range represent
         * weights for the possible values of the distribution.
         */
        template<class Range>
        explicit param_type(const Range& range)
          : _probabilities(boost::begin(range), boost::end(range))
        {
            normalize();
        }

        /**
         * If nw is zero, equivalent to the default constructor.
         * Otherwise, the range of the distribution is [0, nw),
         * and the weights are found by  calling fw with values
         * evenly distributed between \f$\mbox{xmin} + \delta/2\f$ and
         * \f$\mbox{xmax} - \delta/2\f$, where
         * \f$\delta = (\mbox{xmax} - \mbox{xmin})/\mbox{nw}\f$.
         */
        template<class Func>
        param_type(std::size_t nw, double xmin, double xmax, Func fw)
        {
            std::size_t n = (nw == 0) ? 1 : nw;
            double delta = (xmax - xmin) / n;
            BOOST_ASSERT(delta > 0);
            for(std::size_t k = 0; k < n; ++k) {
                _probabilities.push_back(fw(xmin + k*delta + delta/2));
            }
            normalize();
        }

        /**
         * Returns a vector containing the probabilities of each possible
         * value of the distribution.
         */
        std::vector<WeightType> probabilities() const
        {
            return _probabilities;
        }

        /** Writes the parameters to a @c std::ostream. */
        BOOST_RANDOM_DETAIL_OSTREAM_OPERATOR(os, param_type, parm)
        {
            detail::print_vector(os, parm._probabilities);
            return os;
        }
        
        /** Reads the parameters from a @c std::istream. */
        BOOST_RANDOM_DETAIL_ISTREAM_OPERATOR(is, param_type, parm)
        {
            std::vector<WeightType> temp;
            detail::read_vector(is, temp);
            if(is) {
                parm._probabilities.swap(temp);
            }
            return is;
        }

        /** Returns true if the two sets of parameters are the same. */
        BOOST_RANDOM_DETAIL_EQUALITY_OPERATOR(param_type, lhs, rhs)
        {
            return lhs._probabilities == rhs._probabilities;
        }
        /** Returns true if the two sets of parameters are different. */
        BOOST_RANDOM_DETAIL_INEQUALITY_OPERATOR(param_type)
    private:
        /// @cond show_private
        friend class discrete_distribution;
        explicit param_type(const discrete_distribution& dist)
          : _probabilities(dist.probabilities())
        {}
        void normalize()
        {
            impl_type::normalize(_probabilities);
        }
        std::vector<WeightType> _probabilities;
        /// @endcond
    };

    /**
     * Creates a new @c discrete_distribution object that has
     * \f$p(0) = 1\f$ and \f$p(i|i>0) = 0\f$.
     */
    discrete_distribution()
    {
        _impl.init_empty();
    }
    /**
     * Constructs a discrete_distribution from an iterator range.
     * If @c first == @c last, equivalent to the default constructor.
     * Otherwise, the values of the range represent weights for the
     * possible values of the distribution.
     */
    template<class Iter>
    discrete_distribution(Iter first, Iter last)
    {
        init(first, last);
    }
#ifndef BOOST_NO_CXX11_HDR_INITIALIZER_LIST
    /**
     * Constructs a @c discrete_distribution from a @c std::initializer_list.
     * If the @c initializer_list is empty, equivalent to the default
     * constructor.  Otherwise, the values of the @c initializer_list
     * represent weights for the possible values of the distribution.
     * For example, given the distribution
     *
     * @code
     * discrete_distribution<> dist{1, 4, 5};
     * @endcode
     *
     * The probability of a 0 is 1/10, the probability of a 1 is 2/5,
     * the probability of a 2 is 1/2, and no other values are possible.
     */
    discrete_distribution(std::initializer_list<WeightType> wl)
    {
        init(wl.begin(), wl.end());
    }
#endif
    /**
     * Constructs a discrete_distribution from a Boost.Range range.
     * If the range is empty, equivalent to the default constructor.
     * Otherwise, the values of the range represent weights for the
     * possible values of the distribution.
     */
    template<class Range>
    explicit discrete_distribution(const Range& range)
    {
        init(boost::begin(range), boost::end(range));
    }
    /**
     * Constructs a discrete_distribution that approximates a function.
     * If nw is zero, equivalent to the default constructor.
     * Otherwise, the range of the distribution is [0, nw),
     * and the weights are found by  calling fw with values
     * evenly distributed between \f$\mbox{xmin} + \delta/2\f$ and
     * \f$\mbox{xmax} - \delta/2\f$, where
     * \f$\delta = (\mbox{xmax} - \mbox{xmin})/\mbox{nw}\f$.
     */
    template<class Func>
    discrete_distribution(std::size_t nw, double xmin, double xmax, Func fw)
    {
        std::size_t n = (nw == 0) ? 1 : nw;
        double delta = (xmax - xmin) / n;
        BOOST_ASSERT(delta > 0);
        std::vector<WeightType> weights;
        for(std::size_t k = 0; k < n; ++k) {
            weights.push_back(fw(xmin + k*delta + delta/2));
        }
        init(weights.begin(), weights.end());
    }
    /**
     * Constructs a discrete_distribution from its parameters.
     */
    explicit discrete_distribution(const param_type& parm)
    {
        param(parm);
    }

    /**
     * Returns a value distributed according to the parameters of the
     * discrete_distribution.
     */
    template<class URNG>
    IntType operator()(URNG& urng) const
    {
        BOOST_ASSERT(!_impl._alias_table.empty());
        IntType result;
        WeightType test;
        do {
            result = uniform_int_distribution<IntType>((min)(), (max)())(urng);
            test = _impl.test(urng);
        } while(!_impl.accept(result, test));
        if(test < _impl._alias_table[static_cast<std::size_t>(result)].first) {
            return result;
        } else {
            return(_impl._alias_table[static_cast<std::size_t>(result)].second);
        }
    }
    
    /**
     * Returns a value distributed according to the parameters
     * specified by param.
     */
    template<class URNG>
    IntType operator()(URNG& urng, const param_type& parm) const
    {
        if(WeightType limit = impl_type::try_get_sum(parm._probabilities)) {
            WeightType val = impl_type::generate_in_range(urng, limit);
            WeightType sum = 0;
            std::size_t result = 0;
            for(typename std::vector<WeightType>::const_iterator
                    iter = parm._probabilities.begin(),
                    end = parm._probabilities.end();
                iter != end; ++iter, ++result)
            {
                sum += *iter;
                if(sum > val) {
                    return result;
                }
            }
            // This shouldn't be reachable, but round-off error
            // can prevent any match from being found when val is
            // very close to 1.
            return static_cast<IntType>(parm._probabilities.size() - 1);
        } else {
            // WeightType is integral and sum(parm._probabilities)
            // would overflow.  Just use the easy solution.
            return discrete_distribution(parm)(urng);
        }
    }
    
    /** Returns the smallest value that the distribution can produce. */
    result_type min BOOST_PREVENT_MACRO_SUBSTITUTION () const { return 0; }
    /** Returns the largest value that the distribution can produce. */
    result_type max BOOST_PREVENT_MACRO_SUBSTITUTION () const
    { return static_cast<result_type>(_impl._alias_table.size() - 1); }

    /**
     * Returns a vector containing the probabilities of each
     * value of the distribution.  For example, given
     *
     * @code
     * discrete_distribution<> dist = { 1, 4, 5 };
     * std::vector<double> p = dist.param();
     * @endcode
     *
     * the vector, p will contain {0.1, 0.4, 0.5}.
     *
     * If @c WeightType is integral, then the weights
     * will be returned unchanged.
     */
    std::vector<WeightType> probabilities() const
    {
        std::vector<WeightType> result(_impl._alias_table.size(), static_cast<WeightType>(0));
        std::size_t i = 0;
        for(typename impl_type::alias_table_t::const_iterator
                iter = _impl._alias_table.begin(),
                end = _impl._alias_table.end();
                iter != end; ++iter, ++i)
        {
            WeightType val = iter->first;
            result[i] += val;
            result[static_cast<std::size_t>(iter->second)] += _impl.get_weight(i) - val;
        }
        impl_type::normalize(result);
        return(result);
    }

    /** Returns the parameters of the distribution. */
    param_type param() const
    {
        return param_type(*this);
    }
    /** Sets the parameters of the distribution. */
    void param(const param_type& parm)
    {
        init(parm._probabilities.begin(), parm._probabilities.end());
    }
    
    /**
     * Effects: Subsequent uses of the distribution do not depend
     * on values produced by any engine prior to invoking reset.
     */
    void reset() {}

    /** Writes a distribution to a @c std::ostream. */
    BOOST_RANDOM_DETAIL_OSTREAM_OPERATOR(os, discrete_distribution, dd)
    {
        os << dd.param();
        return os;
    }

    /** Reads a distribution from a @c std::istream */
    BOOST_RANDOM_DETAIL_ISTREAM_OPERATOR(is, discrete_distribution, dd)
    {
        param_type parm;
        if(is >> parm) {
            dd.param(parm);
        }
        return is;
    }

    /**
     * Returns true if the two distributions will return the
     * same sequence of values, when passed equal generators.
     */
    BOOST_RANDOM_DETAIL_EQUALITY_OPERATOR(discrete_distribution, lhs, rhs)
    {
        return lhs._impl == rhs._impl;
    }
    /**
     * Returns true if the two distributions may return different
     * sequences of values, when passed equal generators.
     */
    BOOST_RANDOM_DETAIL_INEQUALITY_OPERATOR(discrete_distribution)

private:

    /// @cond show_private

    template<class Iter>
    void init(Iter first, Iter last, std::input_iterator_tag)
    {
        std::vector<WeightType> temp(first, last);
        init(temp.begin(), temp.end());
    }
    template<class Iter>
    void init(Iter first, Iter last, std::forward_iterator_tag)
    {
        std::vector<std::pair<WeightType, IntType> > below_average;
        std::vector<std::pair<WeightType, IntType> > above_average;
        WeightType weight_average = _impl.init_average(first, last);
        WeightType normalized_average = _impl.get_weight(0);
        std::size_t i = 0;
        for(; first != last; ++first, ++i) {
            WeightType val = impl_type::normalize(*first, weight_average);
            std::pair<WeightType, IntType> elem(val, static_cast<IntType>(i));
            if(val < normalized_average) {
                below_average.push_back(elem);
            } else {
                above_average.push_back(elem);
            }
        }

        typename impl_type::alias_table_t::iterator
            b_iter = below_average.begin(),
            b_end = below_average.end(),
            a_iter = above_average.begin(),
            a_end = above_average.end()
            ;
        while(b_iter != b_end && a_iter != a_end) {
            _impl._alias_table[static_cast<std::size_t>(b_iter->second)] =
                std::make_pair(b_iter->first, a_iter->second);
            a_iter->first -= (_impl.get_weight(b_iter->second) - b_iter->first);
            if(a_iter->first < normalized_average) {
                *b_iter = *a_iter++;
            } else {
                ++b_iter;
            }
        }
        for(; b_iter != b_end; ++b_iter) {
            _impl._alias_table[static_cast<std::size_t>(b_iter->second)].first =
                _impl.get_weight(b_iter->second);
        }
        for(; a_iter != a_end; ++a_iter) {
            _impl._alias_table[static_cast<std::size_t>(a_iter->second)].first =
                _impl.get_weight(a_iter->second);
        }
    }
    template<class Iter>
    void init(Iter first, Iter last)
    {
        if(first == last) {
            _impl.init_empty();
        } else {
            typename std::iterator_traits<Iter>::iterator_category category;
            init(first, last, category);
        }
    }
    typedef typename detail::select_alias_table<
        (::boost::is_integral<WeightType>::value)
    >::template apply<IntType, WeightType>::type impl_type;
    impl_type _impl;
    /// @endcond
};

}
}

#include <boost/random/detail/enable_warnings.hpp>

#endif