from sklearn.cluster import MiniBatchKMeans
import numpy as np
import argparse
import cv2
from common import convert_color, closest_color, create_colored_image, remove_from_list, list_match

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

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

	image = None
	
	h = 0
	w = 0
	n_colors = 3

	white = [255, 255, 255]

	output = ''

	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
		
		self.flatten_pallete()
		self.posterize()
		self.determine_colors()

	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 = closest_color(self.original_colors, [255, 255, 255])

		blank = create_colored_image(self.w, self.h, [255, 255, 255])
		composite = create_colored_image(self.w, self.h, [255, 255, 255])
		for i in range(self.n_colors) :
			if list_match(self.original_colors[i], white) :
				continue
			mask = self.extract_color_mask(self.image, self.original_colors[i])
			closest = closest_color(self.colors, self.original_colors[i])
			self.colors = remove_from_list(self.colors, closest)
			name = self.match_color_name(closest)

			cv2.imwrite(f'{name}.png', mask)

			color_mat = create_colored_image(self.w, self.h, closest)
			color_mask = cv2.bitwise_or(color_mat, color_mat, mask = (255-mask))

			


		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]}']