import os
import sys

import cv2 as cv
import numpy as np
import scipy.cluster

from annotations import DataFile,computeBoundingBox


def createHistogram(img):
	# Convert BGR to HSV
	hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV)
	# H in range(0,180)
	# S in range(0,256)
	# V in range(0,256)
	planes=cv.split(hsv)
	hhist=cv.calcHist(planes,[0],None,[256],(0,180),accumulate=False)
	shist=cv.calcHist(planes,[1],None,[256],(0,256),accumulate=False)
	vhist=cv.calcHist(planes,[2],None,[256],(0,256),accumulate=False)

	width=512
	height=400
	binSize=width//256

	histImage = np.zeros((height, width, 3), dtype=np.uint8)
	cv.normalize(hhist, hhist, alpha=0, beta=height, norm_type=cv.NORM_MINMAX)
	cv.normalize(shist, shist, alpha=0, beta=height, norm_type=cv.NORM_MINMAX)
	cv.normalize(vhist, vhist, alpha=0, beta=height, norm_type=cv.NORM_MINMAX)

	for i in range(1, 256):
		cv.line(histImage, ( binSize*(i-1), height - int(round(hhist[i-1][0])) ),
						( binSize*(i), height - int(round(hhist[i][0])) ),
						( 255, 0, 0), thickness=2)
		cv.line(histImage, ( binSize*(i-1), height - int(round(shist[i-1][0])) ),
						( binSize*(i), height - int(round(shist[i][0])) ),
						( 0, 255, 0), thickness=2)
		cv.line(histImage, ( binSize*(i-1), height - int(round(vhist[i-1][0])) ),
						( binSize*(i), height - int(round(vhist[i][0])) ),
						( 0, 0, 255), thickness=2)

	cv.imshow('Source image', img)
	cv.imshow('calcHist Demo', histImage)
	cv.waitKey()


def quantize(img):
	arr=np.reshape(img,(-1,3)).astype(np.float)
	colors=np.array([[0,0,0],[255,255,255],[193,165,116]],np.float)
	print(colors)
	(centers,distortion)=scipy.cluster.vq.kmeans(arr,colors)
	print(centers)
	return centers


def computeClosest(x,centers):
	res=centers[0]
	d=np.linalg.norm(res-x)
	for c in centers:
		d_=np.linalg.norm(c-x)
		if d_<d:
			res=c
			d=d_
	return res


filepath=sys.argv[1]
annotations=DataFile(sys.argv[2])
filename=os.path.basename(filepath)
(x1,y1,x2,y2)=computeBoundingBox(annotations[filename][0])
(w,h)=(x2-x1,y2-y1)
img=cv.imread(filepath)
(x3,x4,y3,y4)=(x1+w//4,x1+3*w//4,y1+h//4,y1+3*h//4)
print(x3,x4,y3,y4)
rect=img[y3:y4,x3:x4,:]
centers=quantize(rect)

# colors=[(0,0,0),(255,255,255),(255,200,0)]
# colorMap=colors[:]
# for (i,c) in enumerate(centers):
# 	colorMap[i]=computeClosest(c,colors)

for x in range(x1,x2):
	for y in range(y1,y2):
		pix=img[y,x]
		img[y,x]=computeClosest(pix,centers)

rect=img[y1:y2,x1:x2]
maskB=cv.inRange(rect,np.array([0,0,0]),np.array([60,60,60]))
maskB=cv.erode(maskB,np.ones((3,3),np.uint8))
maskW=cv.inRange(rect,np.array([160,160,160]),np.array([256,256,256]))
maskW=cv.erode(maskW,np.ones((3,3),np.uint8))

cv.imshow(filename,img)
cv.waitKey()
cv.imshow(filename,maskB)
cv.waitKey()
cv.imshow(filename,maskW)
cv.waitKey()