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

###################################
#
#   BOUNDING BOX
#
"""
A Polytope is defined as the intersection of halfspaces that solves
the inequality `A x <= b`.

Let's do an example:

   [ 1   0]          [10]
   [ 0   1]  [x] <=  [15]
   [-1   0]  [y]     [ 0]
   [ 0  -1]          [ 0]

In other words:

	0 <= x <= 10
	0 <= y <= 15

This example could also have been done with

	box = pc.box2poly([[0, 10], [0, 15]])

"""

def bounding_box(xmin,xmax,ymin,ymax):
	A = np.array([[1,0],[0,1],[-1,-0],[-0,-1]])
	b = np.array([xmax,ymax,-xmin,-ymin])
	return pc.Polytope(A, b)

box = bounding_box(xmin=0,xmax=10,ymin=0,ymax=15)

"""
Some things you can do with polytopes:
"""

print(box)              # pretty-print
box.volume              # 10 * 15 = 150
[2,2] in box            # True
np.array([2,2]) in box  # True

"""
To be clear, we are not restricted to 2D.
"""
U = pc.box2poly([[0, 10], [0, 15], [0,20]])

U.volume     # 10 * 15 * 20 = 3000
[3,4,5] in U # True


"""
We can compute intersections.
"""
U = pc.box2poly([[5, 15], [5, 20]])
V = box.intersect(U)

V.volume # 5 * 10 = 50

"""
Polytope cannot handle intersections with a halfspace.
You have to take intersections of actual polytopes.
"""

A = np.array([[1,1]])
b = np.array([4])
U = pc.Polytope(A, b)

[1,1] in U # You can still do operations.
U.volume   # 0 -- But there is no volume.

print(U.intersect(box)) # And intersections don't work.
print(box.intersect(U)) # And intersections don't work.


###################################
#
#   CANDIDATES
#

"""
Each candidate is a point in space.

Use NumPy arrays so you can do math with them later.
"""

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

###################################
#
#   VORONOI CELLS
#

"""
We need to make the Voronoi cells by hand.
"""

def voronoi_cells(candidates,xmin,xmax,ymin,ymax):
	"""
	Compute the polytopes that correspond to the Voroni cells of
	each candidates, bounded by [xmin,xmax,ymin,ymax].
	"""
	n_candidates = candidates.shape[0]
	
	# Begin one set of inequalities for each candidate.
	A = np.array([[1,0],[0,1],[-1,-0],[-0,-1]])
	b = np.array([xmax,ymax,-xmin,-ymin])
	
	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


###################################
#
#   PLURALITY ELECTION
#

voronoi = voronoi_cells(candidates,xmin=0,xmax=10,ymin=0,ymax=15)

# The cell of the 0th candidate is easy to verify by hand.
print(voronoi[0])  # 0 <= x <= 5 and 0 <= y <= 8
voronoi[0].volume  # 5 * 7 = 35

# Let's compute the 1st place votes.
fp = np.array([v.volume / box.volume for v in voronoi])

print("The Plurality winner is candidate %d with %f percent of the vote." %
 		(fp.argmax(), 100*fp.max())
	)


###################################
#
#   RANKED BALLOTS
#

first_place = voronoi_cells(candidates         ,0,10,0,15) # A,B,C
without_C0  = voronoi_cells(candidates[[1,2],:],0,10,0,15) # B,C
without_C1  = voronoi_cells(candidates[[0,2],:],0,10,0,15) # A,C
without_C2  = voronoi_cells(candidates[[0,1],:],0,10,0,15) # A,B

#
# Let's name the candidates 'A','B','C'
#

ballots = {
	"A>B>C": first_place[0].intersect(without_C0[0]), # A then B
	"A>C>B": first_place[0].intersect(without_C0[1]), # A then C
	"B>A>C": first_place[1].intersect(without_C1[0]), # B then A
	"B>C>A": first_place[1].intersect(without_C1[1]), # B then C
	"C>A>B": first_place[2].intersect(without_C2[0]), # C then A
	"C>B>A": first_place[2].intersect(without_C2[1]), # C then B
}

tally = {
	"A>B>C": ballots["A>B>C"].volume / box.volume,
	"A>C>B": ballots["A>C>B"].volume / box.volume,
	"B>A>C": ballots["B>A>C"].volume / box.volume,
	"B>C>A": ballots["B>C>A"].volume / box.volume,
	"C>A>B": ballots["C>A>B"].volume / box.volume,
	"C>B>A": ballots["C>B>A"].volume / box.volume,
}