import matplotlib.pyplot as plt
import polytope as pc
import numpy as np

###################################
#
#   VORONOI CELLS
#
def voronoi_cells(candidates,bounds):
	"""
	Compute the polytopes that correspond to the Voroni cells of
	each candidates, bounded by the box
	
	      [[x0_min,x0_max],
	       [x1_min,x1_max],
	       ...
	       [xN_min,xN_max]]
	
	where N is the number of dimensions.
	"""
	n_candidates,n_dim = candidates.shape

	assert(n_dim == len(bounds))
	
	# Begin one set of inequalities for each candidate.
	A = np.zeros((2*n_dim,n_dim))
	b = np.zeros(2*n_dim)
	
	for j in range(n_dim):
		k = j + n_dim
		A[j,j] =  1 # [0, 0, ..., 1, 0, ...]
		A[k,j] = -1 # [0, 0, ...,-1, 0, ...]
		b[j] =  bounds[j][1] #  x0_max
		b[k] = -bounds[j][0] # -x0_min
	
	A = [A for j in range(n_candidates)]
	b = [b for j in range(n_candidates)]
	
	for j in range(n_candidates-1):
		for k in range(j+1,n_candidates):
			
			cj = candidates[j,:]
			ck = candidates[k,:]
			
			# Vector normal to the hyperplane.
			v = ck - cj # Vector from cj to ck.
			n = v / np.linalg.norm(v)
			
			# The point half-way between cj and ck is on that plane.
			p = cj + 0.5*v
			q = np.dot(p,n)
			
			# Write `n` as a 2D array.
			n_2d = np.array([[1]]) * n
			
			if np.dot(cj,n) < q:
				# Candidate j:  `n dot p <= q`
				# Candidate k:  `n dot p >= q`  -->  `-n dot p <= -q`
				A[j] = np.append(A[j], n_2d,axis=0)
				A[k] = np.append(A[k],-n_2d,axis=0)
				
				b[j] = np.append(b[j], q)
				b[k] = np.append(b[k],-q)
			else:
				# Candidate j:  `n dot p >= b`  -->  `-n dot p <= -b`
				# Candidate k:  `n dot p <= b`
				A[k] = np.append(A[k], n_2d,axis=0)
				A[j] = np.append(A[j],-n_2d,axis=0)
				
				b[k] = np.append(b[k], q)
				b[j] = np.append(b[j],-q)
				
	voronoi = [pc.Polytope(A[j], b[j]) for j in range(n_candidates)]
	return voronoi

###################################
#
#   COUNTING BALLOTS
#
def compute_ballots(candidates,names,bounds):
	"""
	This function receives a list of candidate positions, their
	names, and the dimensions of the bounding box, and returns
	a dictionary of polytopes corresponding to each possible
	ranked ballot.
	
	USAGE:
	
	ballots = compute_ballots(candidates,['A','B','C'],bounds)
	
	ballots.keys() == ['A>B>C','A>C>B','B>A>C',...]
	"""
	n_candidates = candidates.shape[0]
	first_place  = voronoi_cells(candidates,bounds)
	
	assert(n_candidates == len(names))
	
	#
	# Recursion stops here.
	#
	if n_candidates == 2:
		ballots = {
			names[0]+'>'+names[1]: first_place[0],
			names[1]+'>'+names[0]: first_place[1]
		}
		return ballots
	
	#
	# Recursive part.
	#
	ballots = {}
	for j in range(n_candidates):
		#
		# All candidates except j.
		#
		indices = list(range(n_candidates))
		indices.pop(j)
		
		names_ = names.copy()
		names_.pop(j)
		
		b = compute_ballots(candidates[indices,:],names_,bounds)
		
		for key in b:
			order = names[j] + ">" + key
			ballots[order] = first_place[j].intersect(b[key])
	
	return ballots

def compute_tally(candidates,names,bounds):
	"""
	This function receives a list of candidate positions, their
	names, and the dimensions of the bounding box, and returns
	a dictionary showing the fraction of voters associated with
	each possible ballot.
	
	USAGE:
	
	tally = compute_tally(candidates,['A','B','C'],bounds)
	
	tally.keys() == ['A>B>C','A>C>B','B>A>C',...]
	
	tally['A>B>C'] = fraction of voters that vote A>B>C.
	"""
	ballots = compute_ballots(candidates,names,bounds)
	tally = {}
	box = pc.box2poly(bounds)
	
	for key in ballots:
		tally[key] = ballots[key].volume / box.volume

	return tally


###################################
#
#   MAIN PROGRAM
#

#
# EXAMPLE 1
#
candidates = np.array([
	[2,5],
	[8,5],
	[2,11]
])

bounds = [[0, 10], [0, 15]]
names  = ['A','B','C']
tally  = compute_tally(candidates,names,bounds)


#
# EXAMPLE 2
#
n_candidates = 4
n_dim = 5

candidates = np.random.rand(n_candidates,n_dim)
unit_bound = [ [0,1] for j in range(n_dim) ]

names = ['A','B','C','D']
tally = compute_tally(candidates,names,unit_bound)