import random
from collections import Counter
import itertools

random.seed(0xdeadbeef)
# output for this seed:
# testing good draw
# PASSED chi-squared test; got 2.612083333333333 <  18.47 for 3 DOF
# testing bad draw
# FAILED chi-squared test; got 1005.0618333333333 >= 18.47 for 3 DOF

probs = [
    (1, 0.1),
    (2, 0.2),
    (3, 0.3),
    (4, 0.4),
]

chi_squared_table = [
    # p-values
    [ 0,  0.95,  0.90,   0.80,   0.70,   0.50,    0.30,   0.20,   0.10,   0.05,   0.01,  0.001],
    # values
    [ 1, 0.004,  0.02,   0.06,   0.15,   0.46,    1.07,   1.64,   2.71,   3.84,   6.63,  10.83],
    [ 2,  0.10,  0.21,   0.45,   0.71,   1.39,    2.41,   3.22,   4.61,   5.99,   9.21,  13.82],
    [ 3,  0.35,  0.58,   1.01,   1.42,   2.37,    3.66,   4.64,   6.25,   7.81,  11.34,  16.27],
    [ 4,  0.71,  1.06,   1.65,   2.20,   3.36,    4.88,   5.99,   7.78,   9.49,  13.28,  18.47],
    [ 5,  1.14,  1.61,   2.34,   3.00,   4.35,    6.06,   7.29,   9.24,  11.07,  15.09,  20.52],
    [ 6,  1.63,  2.20,   3.07,   3.83,   5.35,    7.23,   8.56,  10.64,  12.59,  16.81,  22.46],
    [ 7,  2.17,  2.83,   3.82,   4.67,   6.35,    8.38,   9.80,  12.02,  14.07,  18.48,  24.32],
    [ 8,  2.73,  3.49,   4.59,   5.53,   7.34,    9.52,  11.03,  13.36,  15.51,  20.09,  26.12],
    [ 9,  3.32,  4.17,   5.38,   6.39,   8.34,   10.66,  12.24,  14.68,  16.92,  21.67,  27.88],
    [10,  3.94,  4.87,   6.18,   7.27,   9.34,   11.78,  13.44,  15.99,  18.31,  23.21,  29.59],
]

p = 0.001 # controls flakiness; lower value decreases change of random failures
          # but increases chance bad implementation gets through
p_index = chi_squared_table[0].index(p)

def test_draw(draw, num_trials=10000):
    N = len(probs) # N - 1 DOF
    observed = Counter()
    for _ in range(num_trials):
        observed[draw(probs)] += 1
    expected = Counter()
    for val, prob in probs:
        expected[val] = num_trials * prob
    # compute statistical significance
    chi_squared = sum(
        (observed[k] - expected[k]) ** 2 / expected[k]
        for k in set(observed) | set(expected)
    )
    table_value = chi_squared_table[N][p_index]
    if chi_squared >= table_value:
        print(f'FAILED chi-squared test; got {chi_squared} >= {table_value} for {N - 1} DOF')
    else:
        print(f'PASSED chi-squared test; got {chi_squared} <  {table_value} for {N - 1} DOF')

def good_draw(probs):
    rv = random.random()
    for val, prob in probs:
        rv -= prob
        if rv <= 0:
            return val
    # p=0 this happens
    return probs[0][0]

def bad_draw(probs):
    # increasing number of trials makes this converge to good_draw
    for _ in range(2):
        val, prob = random.choice(probs)
        if random.random() < prob:
            return val
    # give up
    return random.choice(probs)[0]

print('testing good draw')
test_draw(good_draw)
print('testing bad draw')
test_draw(bad_draw)