import os
import sys

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

from annotations import DataFile,computeBoundingBox
from hough import show


def kmeans(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("k-means centers:",centers)
	return centers


def quantize(img,centers):
	origShape=img.shape
	data=np.reshape(img,(-1,3))
	(keys,dists)=scipy.cluster.vq.vq(data,centers)
	pixels=np.array([centers[k] for k in keys],dtype=np.uint8).reshape(origShape)
	return pixels


def filterContours(contours,bwImg,stoneDims):
	contourImg=cv.cvtColor(bwImg,cv.COLOR_GRAY2BGR)
	res=[]
	for (i,c) in enumerate(contours):
		keep=True
		print(i)
		moments=cv.moments(c)
		center=(moments["m10"]/(moments["m00"] or 1), moments["m01"]/(moments["m00"] or 1))
		print("center:", center)
		area=cv.contourArea(c)
		print("area:",area)
		(x,y,w,h)=cv.boundingRect(c)
		print("bounding box:",(x,y,w,h))
		if w>stoneDims[0] or h>stoneDims[1]*1.3 or w<2 or h<2:
			cv.drawMarker(contourImg,tuple(map(int,center)),(0,0,255),cv.MARKER_TILTED_CROSS,12)
			keep=False
		coverage1=area/(w*h or 1)
		print("coverage1:",coverage1)
		hull=cv.convexHull(c)
		coverage2=area/(cv.contourArea(hull) or 1)
		print("coverage2:",coverage2)
		if coverage2<0.8:
			cv.drawMarker(contourImg,tuple(map(int,center)),(0,127,255),cv.MARKER_DIAMOND,12)
			keep=False
		print()
		if keep:
			res.append(c)
			cv.drawMarker(contourImg,tuple(map(int,center)),(255,0,0),cv.MARKER_CROSS)
	show(contourImg)
	return res


if __name__=="__main__":
	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:",x3,x4,y3,y4)
	rect=img[y3:y4,x3:x4,:]
	centers=kmeans(rect)
	print("x1,x2,y1,y2:",(x1,x2,y1,y2))
	img[y1:y2,x1:x2,:]=quantize(img[y1:y2,x1:x2,:],centers)
	print("image quantized")

	rect=img[y1:y2,x1:x2]
	unit=np.array([1,1,1],dtype=np.uint8)
	maskB=cv.inRange(rect,centers[0]-unit,centers[0]+unit)
	maskB=cv.dilate(maskB,np.ones((3,3),np.uint8),iterations=1)
	maskB=cv.erode(maskB,np.ones((3,3),np.uint8),iterations=3)
	maskW=cv.inRange(rect,centers[1]-unit,centers[1]+unit)
	maskW=cv.erode(maskW,np.ones((3,3),np.uint8),iterations=2)

	show(img,filename)
	show(maskB,filename)
	show(maskW,filename)
	stones=cv.bitwise_or(maskB,maskW)
	show(stones)

	stoneDims=(w/19,h/19)
	print("stone dims:",tuple(x/2 for x in stoneDims),"-",stoneDims)

	(contours,hierarchy)=cv.findContours(stones,cv.RETR_LIST,cv.CHAIN_APPROX_SIMPLE)
	filterContours(contours,stones,stoneDims)