#!/usr/bin/env python

################################################################################
# ICanHaz3D                                                                    #
################################################################################
# A program to convert between several 3D image formats.                       #
# Author: David J. Oftedal.                                                    #
#                                                                              #
# Depends on Python Imaging Library: http://www.pythonware.com/products/pil/   #
#                                                                              #
# Thanks to:                                                                   #
# vrtifacts.com/how-to-teardowns-tutorials/dump-those-silly-colored-3d-glasses #
# for the method for converting anaglyphs.                                     #
################################################################################

from sys import argv
from sys import exit
from PIL import Image, ImageDraw, ImageEnhance, ImageFilter, ImageStat, ImageChops

################################################################################
# Match the brightness and contrast of one greyscale image to another image.   #
# Used to get comparable brightness levels when separating an anaglyph image.  #
################################################################################
def matchbrightness(theimage, reference):
	if(theimage.mode != "L"):
		theimage = theimage.convert("L")

	if(reference.mode != "L"):
		reference = reference.convert("L")

	refbright = ImageStat.Stat(reference).mean[0]
	refcont = ImageStat.Stat(reference).extrema[0][1] - ImageStat.Stat(reference).extrema[0][0]

# Adjust the brightness of the image.
	newbright = ImageStat.Stat(theimage).mean[0]
	brighten = ImageEnhance.Brightness(theimage)
	theimage = brighten.enhance( float(refbright) / float(newbright) )

# Adjust the contrast.
	newcont = ImageStat.Stat(theimage).extrema[0][1] - ImageStat.Stat(theimage).extrema[0][0]
	fixcontrast = ImageEnhance.Contrast(theimage)
	theimage = fixcontrast.enhance( float(refcont) / float(newcont) )

	return theimage

################################################################################
# Split a red/cyan anaglyph image into two frames.                             #
################################################################################
def splitanaglyph(theimage, type="red-cyan"):
	if(theimage.mode != "RGB"):
		theimage = theimage.convert("RGB")

	# Split the image into channels.
	s = theimage.split()
	# Get the red channel for the left image.
	leftframe = s[1 if type == "green-magenta" else 0]

	# Get green and blue for the right image. We use a sensible conversion
	# to greyscale rather than a completely naive one.
	rightframe = Image.merge("RGB", (s[0], s[1].point(lambda x: 0), s[2]) if type == "green-magenta" else (s[0].point(lambda x: 0), s[1], s[2])).convert("L")

	# Adjust the brightness and contrast of the left and right frame.
	leftframe = matchbrightness(leftframe, theimage)
	rightframe = matchbrightness(rightframe, theimage)

	return (leftframe, rightframe)

################################################################################
# Split a red/cyan anaglyph image into two frames and colour the result.       #
################################################################################
def splitcolouranaglyph(theimage, xblur, yblur, blurThreshold, type="red-cyan"):

	# Get the approximate luminance values by splitting the anaglyph in two.
	(leftframe, rightframe) = splitanaglyph(theimage)

	# Create a scaled-down and blurred colour image.
	# This is more or less a sufficient blur for our purposes.
	origsize = theimage.size
	reducedsize = (origsize[0] / max(int(origsize[0]*xblur*0.01),1), max(int(origsize[1]*yblur*0.01),1))
	colours = theimage.resize(reducedsize, Image.ANTIALIAS)
        colours = colours.filter(ImageFilter.SMOOTH_MORE)
	colours = colours.resize(origsize, Image.ANTIALIAS)

	# Blur each pixel only if it results in a significant change in the red/cyan ratio.
        blurredimage = theimage.copy()
	for y in xrange(0, origsize[1]):
		for x in xrange(0, origsize[0]):
			blurredpixel = colours.getpixel((x, y))
			originalpixel = blurredimage.getpixel((x, y))
			try:
				blurredredratio = float(blurredpixel[1 if type == "green-magenta" else 0]) / float(blurredpixel[0 if type == "green-magenta" else 1] + blurredpixel[2])
				originalredratio = float(originalpixel[1 if type == "green-magenta" else 0]) / float(originalpixel[0 if type == "green-magenta" else 1] + originalpixel[2])
				if ((max(blurredredratio, originalredratio) / min(blurredredratio, originalredratio)) - 1.0) > blurThreshold:
					blurredimage.putpixel((x, y), blurredpixel);
			except:
				blurredimage.putpixel((x, y), blurredpixel);

	# Convert the colour image to the YCbCr colourspace and separate the
	# colours from the luminance.
	blurredimage = blurredimage.convert("YCbCr")
	blurredimage = blurredimage.split()

	# Create new images with the luminance from the original frames and the
	# blurred colours.
	leftframe = Image.merge("YCbCr", (leftframe, blurredimage[1], blurredimage[2]))
	rightframe = Image.merge("YCbCr", (rightframe, blurredimage[1], blurredimage[2]))

	# Convert the images back to RGB.
	leftframe = leftframe.convert("RGB")
	rightframe = rightframe.convert("RGB")

	return (leftframe, rightframe)

################################################################################
# Merge two frames into a red/cyan anaglyph.                                   #
################################################################################
def toanaglyph(leftframe, rightframe, type="red-cyan"):
	if(leftframe.mode != "RGB"):
		leftframe = leftframe.convert("RGB")
	if(rightframe.mode != "RGB"):
		rightframe = rightframe.convert("RGB")

	leftframe = leftframe.split()
	rightframe = rightframe.split()

	return Image.merge("RGB",((rightframe if type == "green-magenta" else leftframe)[0], (leftframe if type == "green-magenta" else rightframe)[1], rightframe[2]))

################################################################################
# Merge two frames into a red/cyan anaglyph of a greyscale image.              #
################################################################################
def togreyscaleanaglyph(leftframe, rightframe):
	if(leftframe.mode != "L"):
		leftframe = leftframe.convert("L")
	if(rightframe.mode != "L"):
		rightframe = rightframe.convert("L")

	return toanaglyph(leftframe, rightframe)

################################################################################
# Split a side-by-side image into two frames.                                  #
################################################################################
def fromsidebyside(theimage):
	imgs = theimage.size

	if imgs[0] % 2 != 0:
		print "This side-by-side image has an odd size!"
		exit(1)

	# Start and end coordinates for the two images.
	i1 = (0, 0, imgs[0]/2, imgs[1])
	i2 = (imgs[0]/2, 0, imgs[0], imgs[1])

	leftframe = theimage.crop(i1)
	rightframe = theimage.crop(i2)

	return (leftframe, rightframe)

################################################################################
# Merge two frames into a side-by-side image.                                  #
################################################################################
def tosidebyside(leftframe, rightframe):
	i=Image.new("RGB",(leftframe.size[0]+rightframe.size[0],max(leftframe.size[1],rightframe.size[1])))
	i.paste(leftframe, (0,0) )
	i.paste(rightframe, (leftframe.size[0],0) )
	return i

################################################################################
# Split a stacked image into two frames.                                       #
################################################################################
def fromstacked(theimage):
	imgs = theimage.size

	if imgs[1] % 2 != 0:
		print "This stacked image has an odd size!"
		exit(1)

	# Start and end coordinates for the two images.
	i1 = (0, 0, imgs[0], imgs[1]/2)
	i2 = (0, imgs[1]/2, imgs[0], imgs[1])

	rightframe = theimage.crop(i1)
	leftframe = theimage.crop(i2)

	return (leftframe, rightframe)

################################################################################
# Merge two frames into a stacked image.                                       #
################################################################################
def tostacked(bottom, top):
	i=Image.new("RGB",(max(top.size[0],bottom.size[0]),top.size[1]+bottom.size[1]))
	i.paste(top, (0,0) )
	i.paste(bottom, (0,top.size[1]) )
	return i

################################################################################
# Split an interlaced image into two frames.                                   #
################################################################################
def frominterlaced(theimage):
	imgs = theimage.size

	# Squeeze the images together and expand them again,
	# dropping unwanted lines and smoothing the result.
	oddframe = theimage.copy()
	evenframe = theimage.copy()

	oddframe = oddframe.resize((imgs[0],imgs[1]/2), Image.NEAREST)
	oddframe = oddframe.resize(imgs, Image.ANTIALIAS)

	#evenframe = evenframe.offset(0, -1) # Drop one line
	evenframe = ImageChops.offset(evenframe, 0, -1) # Drop one line
	evenframe = evenframe.resize((imgs[0],imgs[1]/2), Image.NEAREST)
	evenframe = evenframe.resize(imgs, Image.ANTIALIAS)

	return (evenframe, oddframe)

################################################################################
# Merge two frames into an interlaced image.                                   #
################################################################################
def tointerlaced(evenframe, oddframe):

	if oddframe.size != evenframe.size:
		print "The odd and even images are different sizes!"
		exit(1)

	# Mask to erase every other row in the frames with.
	mask = Image.new("L", evenframe.size);
	draw = ImageDraw.Draw(mask)
	for i in xrange(0, evenframe.size[1], 2):
		draw.line(((0, i), (evenframe.size[0], i)), fill=255)

	return Image.composite(evenframe, oddframe, mask)

################################################################################
# Split a vertically interlaced image into two frames.                         #
################################################################################
def fromcolumns(theimage):
	imgs = theimage.size

	# Squeeze the images together and expand them again,
	# dropping unwanted lines and smoothing the result.
	oddframe = theimage.copy()
	evenframe = theimage.copy()

	oddframe = oddframe.resize((imgs[0]/2,imgs[1]), Image.NEAREST)
	oddframe = oddframe.resize(imgs, Image.ANTIALIAS)

	evenframe = ImageChops.offset(evenframe, -1, 0) # Drop one column
	evenframe = evenframe.resize((imgs[0]/2,imgs[1]), Image.NEAREST)
	evenframe = evenframe.resize(imgs, Image.ANTIALIAS)

	return (evenframe, oddframe)

################################################################################
# Merge two frames into a vertically interlaced image.                         #
################################################################################
def tocolumns(evenframe, oddframe):

	if oddframe.size != evenframe.size:
		print "The odd and even images are different sizes!"
		exit(1)

	# Mask to erase every other column in the frames with.
	mask = Image.new("L", evenframe.size);
	draw = ImageDraw.Draw(mask)
	for i in xrange(0, evenframe.size[0], 2):
		draw.line(((i, 0), (i, evenframe.size[1])), fill=255)

	return Image.composite(evenframe, oddframe, mask)

################################################################################
# Help text.                                                                   #
################################################################################
def usage():
	print
	print "Usage: python " + argv[0] + " [options] image"
	print
	print "The program converts 3D images into other formats."
	print "The conversion is often lossy, so the originals should be kept."
	print
	print "Available options:"
	print "--from format\t\tWhich format to convert from. Default: anaglyph."
	print "--to format\t\tWhich format to convert to. Default: sidebyside."
	print "\tAvailable formats:"
	print "\tanaglyph\tRed-cyan anaglyph."
	print "\tgreen-magenta\tGreen-magenta anaglyph."
	print "\tbwanaglyph\tRed-cyan anaglyph of a greyscale image."
	print "\tinterlaced\tInterlaced rows."
	print "\tcolumns\t\tInterlaced columns."
	print "\tsidebyside\tSide by side."
	print "\tstacked\t\tRight view on top, left on the bottom."
# 	Todo: Legg inn wiggle-greie?
	print "--swap\t\t\tSwap left and right between input and output."
	print "--help\t\t\tShow this help."
	print "--show\t\t\tAttempt to show the resulting image."
	print "--save filename\t\tSave the resulting image in \'filename\'."
	print
	print "Extra options for anaglyph input:"
	print "--xblur num\t\tHow many percent of the width to blur the colour by. Default: 20."
	print "--yblur num\t\tHow many percent of the height to blur the colour by. Default: 20."
	print "--threshold num\t\tHow many percent the red/cyan ratio has to change by in order to have its colour reduced. Default: 50."
	print
	print "Example: python " + argv[0] + " --show image.jpg"
	print

################################################################################
# Entry point to the program.                                                  #
################################################################################
if __name__ == '__main__':
	if len(argv) == 1:
		usage()
		exit(1)

	# Default values:
	fromformat = "anaglyph"
	toformat = "sidebyside"
	xblur = 20
	yblur = 20
	blurThreshold = 0.5
	swap = False
	show = False
	saveimg = False

	# Parse the command line.
	inimage = argv[-1]
	parameterlist = iter(xrange(1, len(argv)-1))
	for i in parameterlist:
		if argv[i] == ("--from"):
			fromformat = argv[i+1]
			next(parameterlist)
		elif argv[i] == ("--to"):
			toformat = argv[i+1]
			next(parameterlist)
		elif argv[i] == ("--xblur"):
			xblur = int(argv[i+1])
			next(parameterlist)
		elif argv[i] == ("--yblur"):
			yblur = int(argv[i+1])
			next(parameterlist)
		elif argv[i] == ("--threshold"):
			blurThreshold = float(argv[i+1]) / 100.0
			next(parameterlist)
		elif argv[i] == "--swap":
			swap = True
		elif argv[i] == "--help":
			usage()
			exit(0)
		elif argv[i] == "--show":
			show = True
		elif argv[i] == ("--save"):
			saveimg = argv[i+1]
			next(parameterlist)
		else:
			print "The option " + argv[i] + " wasn\'t understood."
			exit(1)

	try:
		imagefile = Image.open(inimage)
		imagefile.load()
	except:
		print "The image " + inimage + " couldn\'t be opened."
		exit(1)

	# Input the image.
	if fromformat == "anaglyph":
		(leftframe, rightframe) = splitcolouranaglyph(imagefile, xblur, yblur, blurThreshold)
	elif fromformat == "green-magenta":
		(leftframe, rightframe) = splitcolouranaglyph(imagefile, xblur, yblur, blurThreshold, "green-magenta")
	elif fromformat == "bwanaglyph":
		(leftframe, rightframe) = splitanaglyph(imagefile)
	elif fromformat == "columns":
		(leftframe, rightframe) = fromcolumns(imagefile)
	elif fromformat == "interlaced":
		(leftframe, rightframe) = frominterlaced(imagefile)
	elif fromformat == "sidebyside":
		(leftframe, rightframe) = fromsidebyside(imagefile)
	elif fromformat == "stacked":
		(leftframe, rightframe) = fromstacked(imagefile)
	else:
		print "The input format " + fromformat + " is unknown."
		exit(1)

	if swap == True:
		tmp = leftframe
		leftframe = rightframe
		rightframe = tmp

	# Output the image.
	if toformat == "anaglyph":
		finalimage = toanaglyph(leftframe, rightframe)
	elif toformat == "green-magenta":
		finalimage = toanaglyph(leftframe, rightframe, "green-magenta")
	elif toformat == "bwanaglyph":
		finalimage = togreyscaleanaglyph(leftframe, rightframe)
	elif toformat == "columns":
		finalimage = tocolumns(leftframe, rightframe)
	elif toformat == "interlaced":
		finalimage = tointerlaced(leftframe, rightframe)
	elif toformat == "sidebyside":
		finalimage = tosidebyside(leftframe, rightframe)
	elif toformat == "stacked":
		finalimage = tostacked(leftframe, rightframe)
	else:
		print "The output format " + toformat + " is unknown."
		exit(1)

	if saveimg != False:
		try:
			finalimage.save(saveimg)
		except:
			print "The image could not be saved as " + saveimg + ". Please check the file extension and location."

	if show == True:
		finalimage.show()