#!/usr/bin/env python

# Instagram Unshredder - Binary Simplifier
# by Jason A. Donenfeld <Jason@zx2c4.com>
# for http://instagram-engineering.tumblr.com/post/12651721845/instagram-engineering-challenge-the-unshredder

from PIL import Image
from sys import argv, exit, stderr
from math import sqrt

if len(argv) < 4:
	print >> stderr, "Usage: %s columns input output" % argv[0]
	exit(-1)

image = Image.open(argv[2])
data = image.getdata()
columns = int(argv[1])
shred_width = image.size[0] / columns

def pixel(x, y):
	pixel = data[y * image.size[0] + x]
	return pixel
def pixel_distance(a, b):
	opacity_a = a[3] / 255
	opacity_b = b[3] / 255
	a_r, a_g, a_b = (a[0] * opacity_a, a[1] * opacity_a, a[2] * opacity_a)
	b_r, b_g, b_b = (b[0] * opacity_b, b[1] * opacity_b, b[2] * opacity_b)
	return sqrt((a_r - b_r) ** 2 + (a_g - b_g) ** 2 + (a_b - b_b) ** 2)
def match_score(left, right):
	left = left[len(left) - 1]
	right = right[0]
	straight = 0
	for i in xrange(0, image.size[1]):
		straight += pixel_distance(pixel((left + 1) * shred_width - 1, i), pixel(right * shred_width, i))
	shifted = 0
	offset = image.size[1] / 35
	for i in xrange(offset, image.size[1] - offset):
		shifted += pixel_distance(pixel((left + 1) * shred_width - 1, i + offset), pixel(right * shred_width, i - offset))
	return straight / shifted
def binary_simplify(order_lists):
	simplified = list()
	blacklist = list()
	for i in order_lists:
		if i in blacklist:
			continue;
		minimum = image.size[1] * 442 # = sqrt(3) * 255
		value = None
		right = False
		for j in order_lists:
			if i == j or j in blacklist:
				continue
			score = match_score(i, j)
			if score < minimum:
				minimum = score
				value = j
				right = True
			score = match_score(j, i)
			if score < minimum:
				minimum = score
				value = j
				right = False
		if value == None:
			smushed = i[:]
		elif right:
			smushed = i[:]
			smushed.extend(value)
		else:
			smushed = value[:]
			smushed.extend(i)
		simplified.append(smushed)
		blacklist.append(i)
		blacklist.append(value)
	return simplified

column_order = [[i] for i in xrange(0, columns)]
print column_order
while len(column_order) > 1:
	column_order = binary_simplify(column_order)
	print column_order
column_order = column_order[0]
	
unshredded = Image.new("RGBA", image.size)
for destination, source in enumerate(column_order):
	unshredded.paste(image.crop((shred_width * source, 0, shred_width * source + shred_width, unshredded.size[1])), (destination * shred_width, 0))
unshredded.save(argv[3])