from sklearn.cluster import MiniBatchKMeans
import numpy as np
import argparse
import cv2
from common import convert_color, closest_color_weighted_euclidean, closest_color_euclidean, create_colored_image, remove_from_list, list_match
import os
import subprocess

class Posterize:
	"""Posterize an image and then find nearest colors to use"""
	colors = []
	colors_dict = {}
	original_colors = []
	layers = []
	previews = []

	pallete = None
	pallete_space = 'BGR'
	comparison_space = 'BGR'

	image = None
	
	h = 0
	w = 0
	n_colors = 3
	max_particles = 3000
	conf = os.path.abspath('./conf/base.conf')
	stipple_gen = os.path.abspath('../../../src/stipple_gen')

	white = [255, 255, 255]

	output = None

	def __init__ (self, image, pallete, n_colors, output) :
		self.image = cv2.imread(image)
		(self.h, self.w) = self.image.shape[:2]
		self.pallete = pallete
		self.n_colors = n_colors + 1
		self.output = output

		if not os.path.exists(self.output) :
			print(f'Output directory {self.output} does not exist, creating...')
			os.makedirs(self.output)
		
		self.flatten_pallete()
		self.posterize()
		self.determine_colors()
		self.stipple()
		self.preview()

	def posterize (self):
		lab = cv2.cvtColor(self.image, cv2.COLOR_BGR2LAB)
		feature = lab.reshape((self.h * self.w, 3))
		
		clusters = MiniBatchKMeans(n_clusters = self.n_colors, n_init = 'auto')
		
		labels = clusters.fit_predict(feature)
		quant = clusters.cluster_centers_.astype('uint8')[labels]
		
		rquant = quant.reshape((self.h, self.w, 3))
		rfeature = feature.reshape((self.h, self.w, 3))

		bgrquant = cv2.cvtColor(rquant, cv2.COLOR_LAB2BGR)
		#bgrfeature = cv2.cvtColor(rfeature, cv2.COLOR_LAB2BGR)

		self.image = bgrquant
		
		cv2.imshow('image', bgrquant)
		cv2.waitKey(0)
		cv2.destroyAllWindows()

	def determine_colors (self):
		reshaped = self.image.reshape(-1, self.image.shape[2])
		self.original_colors = np.unique(reshaped, axis=0)
		white, white_dist = closest_color_weighted_euclidean(self.original_colors, [255, 255, 255], 'BGR')

		blank = create_colored_image(self.w, self.h, [255, 255, 255])
		composite = create_colored_image(self.w, self.h, [255, 255, 255])
		
		mask = self.extract_color_mask(self.image, white)
		layer_name = f'WHITE.png'
		output_layer = os.path.join(self.output, layer_name)
		cv2.imwrite(output_layer, mask)
		self.layers.append({
			'layer' : output_layer,
			'color' : white,
			'space' : self.pallete_space
		})

		for i in range(self.n_colors) :
			if list_match(self.original_colors[i], white) :
				continue
			original = self.original_colors[i]
			mask = self.extract_color_mask(self.image, original)
			original_normalized = convert_color(original, 'BGR', self.pallete_space)
			if self.pallete_space == 'RGB' or self.pallete_space == 'BGR' :
				closest, dist = closest_color_weighted_euclidean(self.colors, original_normalized, self.pallete_space)
			else :
				closest, dist = closest_color_euclidean(self.colors, original_normalized)
			self.colors = remove_from_list(self.colors, closest)
			name = self.match_color_name(closest)
			layer_name = f'{name}.png'
			output_layer = os.path.join(self.output, layer_name)
			cv2.imwrite(output_layer, mask)
			self.layers.append({
				'layer' : output_layer,
				'color' : closest,
				'space' : self.pallete_space
			})
			mask = cv2.bitwise_not(mask)
			composite[mask > 0] = np.array(closest)

		cv2.imshow('image', composite)
		cv2.waitKey(0)
		cv2.destroyAllWindows()

	def extract_color_mask (self, image, color):
		mask = cv2.inRange(image, color, color)
		return cv2.bitwise_not(mask)

	def flatten_pallete (self) :
		for color in self.pallete.colors:
			self.colors.append(color['color'])
			self.colors_dict[f'{color["color"][0]},{color["color"][1]},{color["color"][2]}'] = color['name']
			self.pallete_space = color['space']

	def match_color_name (self, key) :
		return self.colors_dict[f'{key[0]},{key[1]},{key[2]}']

	def stipple (self) :
		sanity_check = 0
		for layer in self.layers :
			if 'WHITE.png' in layer['layer'] :
				continue
			l = cv2.imread(layer['layer'], 0)
			(h, w) = l.shape[:2]
			total = h * w
			black = total - cv2.countNonZero(l)
			ratio = black/total
			max_particles = round(ratio * self.max_particles)
			input_image = os.path.abspath(layer['layer'])
			dir_name, file_name = os.path.split(input_image)
			file_part, ext = os.path.splitext(file_name)
			output_image = os.path.join(dir_name, f'{file_part}_preview.png')
			output_svg = os.path.join(dir_name, f'{file_part}.svg')
			cmd = [
				'bash',
				'stipple_gen.sh',
				'--inputImage',  input_image,
				'--outputImage', output_image,
				'--outputSVG', output_svg,
				'--config', self.conf, 
				'--maxParticles', str(max_particles) 
			]
			print(cmd)
			subprocess.call(cmd, cwd = self.stipple_gen)
			self.previews.append(output_image)

	def preview (self) :
		#composite = create_colored_image(self.w, self.h, [255, 255, 255])
		print(self.previews)