from scipy import misc

IMAGE = '1.jpeg'

T = [194, 0, 11]
S = [64, 49, 46]

def dist(x, y):
	res = 0
	for i in xrange(0,2):
		res += (int(x[i]) - int(y[i]))**2
	return res

def get_neighbours(i, j, max_i, max_j):
	result = []
	if i > 0:
		result.append((i - 1, j))
	if j > 0:
		result.append((i, j - 1))
	if (i + 1) < max_i:
		result.append((i + 1, j))
	if (j + 1) < max_j:
		result.append((i, j  + 1))

	return result

def flip(i, j, A, B):
	if (i, j) in A:
		A.remove((i, j))
		B.add((i, j))
	else:
		B.remove((i, j))
		A.add((i, j))

def cut(image):
	dim_x, dim_y, colors = image.shape
	A = set()
	B = {(i, j) for i in range(dim_x) for j in range(dim_y)}
	for i in xrange(0, dim_x):
		for j in xrange(0, dim_y):
			in_A = dist(image[i, j], S)
			in_B = dist(image[i, j], T)
			if in_A > in_B:
				if (i, j) in A:
					flip(i, j, A, B)
			else:
				if (i, j) in B:
					flip(i, j, A, B)


	return A, B


if __name__ == "__main__":
	image = misc.imread(IMAGE)
	dim_x, dim_y, colors = image.shape
	print "Loaded image of shape {x}, {y}".format(x=dim_x, y=dim_y)

	A, B = cut(image)

	for i in range(dim_x):
		for j in range(dim_y):
			if (i, j) in A:
				image[i, j, 0] = min(255, image[i, j, 0] * 2 + 100)
				image[i, j, 1] = image[i, j, 1] / 3
				image[i, j, 2] = image[i, j, 2] / 3
			elif (i,j) in B:
				image[i, j, 0] = image[i, j, 0] / 3
				image[i, j, 1] = min(255, image[i, j, 1] * 2 + 100)
				image[i, j, 2] = image[i, j, 2] / 3

	misc.imsave('out.png', image)