#! /usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Does mosiacs.
"""
import os
import time
import colorsys
# import threading
from multiprocessing import Pool
from multiprocessing.dummy import Pool as ThreadPool

import numpy as np
from PIL import Image

class Mosiac():
	def __init__(self):
		self.numClusters = 8
		self.tolerance = 0.1
		self.numThreads = 4

		self.imageBank = {}
		self.tilesDir = None
		self.clusters = []
		self.clusterMeans = []

		self.bigImageSize = None
		self.bigImage = None
		self.smallImage = None
		self.imageMatrix = []
		self.tileSize = None
		self.matrixSize = None


	def openBigImage(self, bigImagePath, mat=None):
		"""
		Opens and initializes the big image.
		"""
		self.bigImage = Image.open(bigImagePath).convert("RGB")
		# TODO: find a better way to do this
		if not mat:
			if self.bigImage.size[0] > self.bigImage.size[1]:
				mat = (40, 20)
			elif self.bigImage.size[0] < self.bigImage.size[1]:
				mat = (20, 40)
			else:
				mat = (20, 20)

		self.matrixSize = mat
		self.tileSize = (
			self.bigImage.size[0] // mat[0],
			self.bigImage.size[1] // mat[1])
		self.smallImage = self.bigImage.resize(mat)
		self.bigImageSize = (self.tileSize[0]*mat[0], self.tileSize[1]*mat[1])
		self.bigImage = self.bigImage.crop(
			(0, 0, self.bigImageSize[0], self.bigImageSize[1]))


	def initImageBank(self, tilesDir):
		"""
		Calculates the average pixel value of all the images in the directory.
		This is the thread controller.
		"""
		then = time.time()
		self.tilesDir = tilesDir
		files = os.listdir(self.tilesDir)

		pool = ThreadPool()
		pool.map(self.initImageBankWorker, files)
		pool.close()
		pool.join()
		print(f"initImageBank took: {time.time()-then} seconds.")


	def initImageBankWorker(self, file):
		"""
		Thread worker.
		"""
		image = Image.open(os.path.join(self.tilesDir, file))

		if image.mode == "P":
			image = image.convert(image.palette.mode)
		if image.mode == "RGBA":
			image = alpha_composite(image).convert("RGB")
		if image.mode == "L":
			image = image.convert("RGB")

		image = image.resize(self.tileSize, Image.ANTIALIAS)
		mean = tuple(np.mean(image, axis=(0,1)))

		self.imageBank[mean] = image


	def debugImageBank(self):
		"""
		Create a bank of test images.
		"""
		then = time.time()
		for i in range(0, 256, 15):
			for j in range(0, 256, 15):
				for k in range(0, 256, 15):
					image = Image.new("RGB", self.tileSize, (i,j,k, 255))
					mean = tuple(np.mean(image, axis=(0,1)))
					self.imageBank[mean] = image
		print(f"debugImageBank took: {time.time()-then} seconds.")


	def initClusters(self):
		"""
		Divides the images into clusters.
		"""
		then = time.time()
		sort = sorted(self.imageBank.keys(), key=lambda pixel: step(*pixel,8))
		num = len(sort) // self.numClusters

		for n in range(self.numClusters):
			cluster = {pixel: self.imageBank[pixel] for pixel in \
				sort[n*num:n*num+num]}
			self.clusters.append(cluster)

		# shove the left overs into the last cluster
		self.clusters[-1].update({pixel: self.imageBank[pixel] for pixel in \
			sort[-(len(sort) % num):]})

		for cluster in self.clusters:
			mean = tuple(np.mean(list(cluster.keys()), axis=(0)))
			self.clusterMeans.append(mean)
		print(f"initClusters took: {time.time()-then} seconds.")


	def nearestImage(self, pixel):
		"""
		Finds the nearest image within the mosiac image bank to the provided
		pixel.
		"""
		dists = []
		for mean in self.clusterMeans:
			dist = np.linalg.norm(np.array(pixel)-np.array(mean))
			dists.append(dist)
		clusterMean = self.clusterMeans[dists.index(min(dists))]
		cluster = self.clusters[self.clusterMeans.index(clusterMean)]

		nodes = np.array(list(cluster.keys()))
		dist = np.sum((nodes - np.array(pixel))**2, axis=1)
		tol = int(len(dist) * self.tolerance)
		tol += int(tol == 0)
		indexs = dist.argsort()[:tol]
		# choice = indexs.argmin()  # closet value
		choice = np.random.choice(indexs)
		return cluster[tuple(nodes[choice])]


	def buildMatrix(self, tileAlpha):
		"""
		Build the image matrix.
		"""
		then = time.time()
		pixels = list(self.smallImage.getdata())
		for pixel in pixels:
			image = self.nearestImage(pixel)
			if tileAlpha:
				image = image.convert("RGBA")
				comp = Image.new("RGBA", image.size, pixel + (tileAlpha,))
				image = Image.alpha_composite(image, comp).convert("RGB")
			self.imageMatrix.append(image)
		print(f"buildMatrix took: {time.time()-then} seconds.")


	def buildMosiac(self, output, bigAlpha):
		"""
		Builds the final mosiac image.
		"""
		then = time.time()
		image = Image.new("RGB", self.bigImageSize, (255,255,255))

		n = 0
		for y in range(0, self.bigImageSize[1], self.tileSize[1]):
			for x in range(0, self.bigImageSize[0], self.tileSize[0]):
				image.paste(self.imageMatrix[n], box=(x,y))
				n += 1
		if bigAlpha:
			image = image.convert("RGBA")
			self.bigImage.putalpha(bigAlpha)
			image = Image.alpha_composite(image, self.bigImage).convert("RGB")
		# self.bigImage.save(output, "PNG")
		image.save(output, "JPEG", optimize=True, quality=90)
		print(f"buildMosiac took: {time.time()-then} seconds.")


def step (r,g,b, repetitions=1):
	lum = ( .241 * r + .691 * g + .068 * b )**0.5

	h, s, v = colorsys.rgb_to_hsv(r,g,b)

	h2 = int(h * repetitions)
	lum2 = int(lum * repetitions)
	v2 = int(v * repetitions)

	return (h2, lum, v2)


def alpha_composite(image, color=(255, 255, 255)):
	"""
	Alpha composite an RGBA Image with a specified color.
	Source: http://stackoverflow.com/a/9166671/284318
	"""
	image.load()  # needed for split()
	background = Image.new('RGB', image.size, color)
	background.paste(image, mask=image.split()[3])  # 3 is the alpha channel
	return background


if __name__ == "__main__":
	import argparse

	parser = argparse.ArgumentParser(
		description="Stiches together a series of smaller images in the \
		likeliness of a larger image. The 'big image' should be quite large.")
	parser.add_argument(
		"bigImagePath",
		help="The big image that will be used as the 'guide'.")
	parser.add_argument(
		"tilesDir",
		help="Directory full of images to be used as the tiles.")
	parser.add_argument(
		"output",
		help="File to be outputed.")
	parser.add_argument(
		"--matrix-size",
		dest="matrixSize",
		nargs=2,
		type=int,
		help="Size of the tile matrix. A 40x20 matrix would be '40 20'")
	parser.add_argument(
		"--tile-alpha",
		dest="tileAlpha",
		default=50,
		help="Alpha channel value of the color filter that is applied to each \
		tile. Range should be 0-255.")
	parser.add_argument(
		"--big-alpha",
		dest="bigAlpha",
		default=50,
		help="Alpha channel value of big image when it is transposed onto the \
		matrix. Range should be 0-255")
	args = parser.parse_args()

	if args.matrixSize:
		args.matrixSize = tuple(args.matrixSize)

	mosiac = Mosiac()
	mosiac.openBigImage(args.bigImagePath, args.matrixSize)
	mosiac.initImageBank(args.tilesDir)
	# mosiac.debugImageBank()
	mosiac.initClusters()
	mosiac.buildMatrix(args.tileAlpha)
	mosiac.buildMosiac(args.output, args.bigAlpha)