import numpy as np
from numpy.random import uniform


def softmax(values: np.array, eps=1.0E-6) -> np.array:
	result = (values + eps) / (sum(values + eps))
	return result


def sample(values: np.array, with_debug=False):
	# Normalize and sort
	probs = softmax(values)
	sorted_indices = np.argsort(probs)
	sorted_probs = probs[sorted_indices]
	z = uniform()
	if with_debug:
		print(f"Probs={probs}")
		print(f"Z={z:.3f}")
	index = None
	p_sum = 0
	for idx, p in enumerate(sorted_probs):
		p_sum += p
		if z <= p_sum:
			index = sorted_indices[idx]
			break

	return index


def test_sample(data):
	p = np.array(data)
	c = np.array([0]*len(data))

	for i in range(0, 1000):
		index = sample(p, with_debug=False)
		c[index] += 1

	print(c)


if __name__ == '__main__':
	print("Testing sample()")
	test_sample([0.7, 0.1, 0.1, 0.1])
	test_sample([0.1, 0.1, 0.2, 0.1])
	test_sample([0.2, 0.8])
	test_sample([0.1, 0.1])
	test_sample([0.3, 0.0, 0.2, 0.0])
	test_sample([0.0, 0.0, 0.0, 0.0])
	test_sample([0.5, 0.0, 0.2, 0.3])
	test_sample([0.3, 0.1, 0.4, 0.2])