#!/usr/bin/env python

# Instagram Unshredder
# 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) < 3:
	print >> stderr, "Usage: %s input output" % argv[0]
	exit(-1)

image = Image.open(argv[1])
data = image.getdata()

def pixel(x, y):
	return data[y * image.size[0] + x]
def pixel_distance(a, b):
	try:
		opacity_a = a[3] / 255
		opacity_b = b[3] / 255
	except:
		opacity_a = 1
		opacity_b = 1
	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)
columns = []
for col in xrange(0, image.size[0] - 1):
	columns.append(0)
	for row in xrange(0, image.size[1]):
		columns[col] += pixel_distance(pixel(col, row), pixel(col + 1, row))
shred_width = -1
for factor in xrange(2, image.size[0] / 2):
	if image.size[0] % factor != 0:	# Poor man's factorization! R.I.P. Sieve. Sorry instagram.
		continue
	last_col = 0
	success = True
	for col in xrange(0, image.size[0], factor):
		if col == 0:
			continue
		success = True
		for prior in xrange(last_col, col - 1):
			if columns[prior] > columns[col - 1]:
				success = False
				break
		last_col = col
		if not success:
			break
	if success:
		shred_width = factor
		break
if shred_width == -1:
	print >> stderr, "Error: Couldn't determine shred width."
	exit(-2)

columns = image.size[0] / shred_width
print "%s image has %d columns." % (argv[1], columns)

def match_score(left, right):
	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

column_matches = {}
for i in xrange(0, columns):
	min_score = image.size[1] * 442 # = sqrt(3) * 255
	min_index = -1
	for j in xrange(0, columns):
		if j == i:
			continue;
		score = match_score(i, j)
		if score < min_score:
			min_score = score
			min_index = j
	column_matches[i] = (min_index, min_score)

last = -1
# If the graph has any nodes that point to prior ones, let's find the last one this way.
for key, value in column_matches.iteritems():
	for key2, value2 in column_matches.iteritems():
		if key == key2:
			continue
		if value[0] == value2[0]:
			if value[1] > value2[1]:
				last = key
			else:
				last = key2
			break
# Otherwise, if it amazingly wraps around, let's just remove the weakest link.
if last == -1:
	max_score = 0
	for key, value in column_matches.iteritems():
		if value[1] > max_score:
			max_score = value[1]
			last = value[0]
del column_matches[last]

column_matches_reverse = {}
for key, value in column_matches.iteritems():
	column_matches_reverse[value[0]] = (key, value[1])

column_order = []
i = last
while True:
	column_order.append(i)
	if i not in column_matches_reverse:
		break
	i = column_matches_reverse[i][0]
column_order.reverse()

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[2])