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

class Posterize:
	"""Posterize an image and then find nearest colors to use"""
	colors = []
	original_colors = []
	image = None
	pallete = None
	h = 0
	w = 0
	n_colors = 3
	white = [255, 255, 255]

	def __init__ (self, image, pallete, n_colors) :
		self.image = cv2.imread(image)
		(self.h, self.w) = self.image.shape[:2]
		self.pallete = pallete
		self.n_colors = n_colors + 1
		
		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)
		#print(self.original_colors)
		for i in range(self.n_colors) :
			mask = self.extract_color_mask(self.image, self.original_colors[i])
			cv2.imwrite(f'{i}.png', mask)

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