import bisect
import random
import unittest

try:
    xrange
except NameError:
    # Python 3.x
    xrange = range


def weighted_random_choice(seq, weight):
    """Returns a random element from ``seq``. The probability for each element
    ``elem`` in ``seq`` to be selected is weighted by ``weight(elem)``.

    ``seq`` must be an iterable containing more than one element.

    ``weight`` must be a callable accepting one argument, and returning a
    non-negative number. If ``weight(elem)`` is zero, ``elem`` will not be
    considered. 
        
    """ 
    weights = 0
    elems = [] 
    for elem in seq:
        w = weight(elem)     
        try:
            is_neg = w < 0
        except TypeError:    
            raise ValueError("Weight of element '%s' is not a number (%s)" %
                             (elem, w))
        if is_neg:
            raise ValueError("Weight of element '%s' is negative (%s)" %
                             (elem, w))
        if w != 0:               
            try:
                weights += w
            except TypeError:
                raise ValueError("Weight of element '%s' is not a number "
                                 "(%s)" % (elem, w))
            elems.append((weights, elem))
    if not elems:
        raise ValueError("Empty sequence")
    ix = bisect.bisect(elems, (random.uniform(0, weights), None))
    return elems[ix][1]
        

class TestCase(unittest.TestCase):

    def test_empty(self):
        """Empty sequences raise ``ValueError``.

        """
        self.assertRaises(ValueError,
                          weighted_random_choice, [], lambda x: 0)
        self.assertRaises(ValueError,
                          weighted_random_choice, [1, 2, 3], lambda x: 0)

    def test_invalid_weight(self):
        """Invalid weight values are detected.

        """
        self.assertRaises(ValueError,
                          weighted_random_choice, [1, 2, 3], lambda x: "foo")

        class Oops(Exception):
            pass

        def weight(elem):
            raise Oops()

        self.assertRaises(Oops, weighted_random_choice, [1, 2, 3], weight)

    def test_spread(self):
        """Results are consistent with weight function.

        """
        seq = range(0, 100)
        odds, evens = [], []

        bias = 10.0

        def weight(elem):
            if elem % 2:
                return bias
            else:
                return 1

        for _ in xrange(0, 5000):
            elem = weighted_random_choice(seq, weight)
            if elem % 2:
                odds.append(elem)
            else:
                evens.append(elem)

        delta = abs(bias - float(len(odds) / float(len(evens))))
        self.assertTrue(delta < 1)


if __name__ == "__main__":
    random.seed()
    unittest.main()