gistfile1.py

'''
Id:          "$Id$"
Description: 
Test:        
'''

"""
Given a list of probabilities (floats), and a list of values (anything), write a function/class/whatever:

go(probs, values)

that returns the values with the given probabilities.
For example, you'd expect calling:

go([0.3, 0.7], ['hello', 'world'])

100 times would give (roughly, but probably not exactly) 30 'hello's and 70 'world's, in some order.

Bear in mind that you don't know in advance how many times 'go' is going to be called, so it's basically a random value generator.

There are many ways to write this, many interfaces -- go could be a function, a class, whatever you want.
Try to do it in the most Pythonic way possible, and write some unit tests.

"""

import random
import unittest2
#import itertools
#from nose.tools import eq_, raises
#from collections import defaultdict
#from contextlib import contextmanager
from qz.core.test.qzunittest import qzUnitTest    

from exceptions import Exception

class ProbTestInitException(Exception):
    pass

class ProbTest(object):
    
    def __init__(self, probs, values):
        """
        We expect two equal-sized lists of probabilities and values.
        """
        if len(probs) != len(values):
            raise ProbTestInitException
        # if sum(probs) != 1.0:
            # raise ProbTestInitException
        combined = zip(probs, values)
        
        self.pt = {}
        for p, v in combined:
            self.pt[p] = v
            
    def go(self):
        r = random.random()
        for p, v in self.pt.items():
            if r <= p:
                return self.pt[p]

class TestGo(unittest2.TestCase):
    
    def test_pass(self):
        pass
        
    def test_failure_unless_both_len_of_probs_and_values_are_equal(self):
        probs = [0.1,]
        values = ['foo', 'bar,']
        with self.assertRaises(ProbTestInitException):
            ProbTest(probs, values)
        
    def test_failure_unless_probs_values_sum_to_one(self):
        probs = [0.1, 0.4, 0.2,]
        values = ['foo', 'bar', 'baz', ]
        with self.assertRaises(ProbTestInitException):
            ProbTest(probs, values)

    def test_successful_init(self):
        probs = [0.4, 0.6,]
        values = ['foo', 'bar', ]
        pt = ProbTest(probs, values)
        self.assertEqual(isinstance(pt, ProbTest), True)

    def test_always_expected_value(self):
        probs = [1.0,]
        values = ['foo', ]
        pt = ProbTest(probs, values)
        self.assertEqual(pt.go(), 'foo')
        
    def test_values(self):
        probs = [0.4, 0.6,]
        values = ['foo', 'bar', ]
        pt = ProbTest(probs, values)
        print pt.go()

            
def main():    
    #qzUnitTest(headless=True, suppressOutput=False)
    
    probs = [0.1, 0.2, 0.4, 0.3,]
    probs_actual = [0.1, 0.3, 0.7, 1.0,]
    values = ['one', 'two', 'three', 'four', ]
    pt = ProbTest(probs_actual, values)
    count_dict = {}
    for i in range(1000):
        key = pt.go()
        count_dict[key] = count_dict.setdefault(key, 1) + 1
    print count_dict