import sys
sys.path.append("../src")

import os
import math
import random

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

from annotations import DataFile,computeBoundingBox
from hough import show
from analyzer.epoint import EPoint

random.seed(361)


class Line():
	def __init__(self,a,b):
		self.a=a
		self.b=b
		self.points={a,b}

	def getSortedPoints(self):
		return tuple(sorted(self.points))


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 filterStones(contours,bwImg,stoneDims):
	contourImg=cv.cvtColor(bwImg,cv.COLOR_GRAY2BGR)
	res=[]
	for (i,c) in enumerate(contours):
		keep=True
		moments=cv.moments(c)
		center=(moments["m10"]/(moments["m00"] or 1), moments["m01"]/(moments["m00"] or 1))
		area=cv.contourArea(c)
		(x,y,w,h)=cv.boundingRect(c)
		if w>stoneDims[0] or h>stoneDims[1]*1.5 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)
		hull=cv.convexHull(c)
		coverage2=area/(cv.contourArea(hull) or 1)
		if coverage2<0.8:
			cv.drawMarker(contourImg,tuple(map(int,center)),(0,127,255),cv.MARKER_DIAMOND,12)
			keep=False
		if keep:
			res.append((EPoint(*center),c))
			cv.drawMarker(contourImg,tuple(map(int,center)),(255,0,0),cv.MARKER_CROSS)
	print("accepted:",len(res))
	print("rejected:",len(contours)-len(res))
	show(contourImg)
	return res


def point2lineDistance(a,b,p):
	# https://en.wikipedia.org/wiki/Point-line_distance#Line_defined_by_two_points
	ab=b-a
	num=abs(ab.y*p.x - ab.x*p.y + b.x*a.y - a.x*b.y)
	denum=math.sqrt(ab.y**2+ab.x**2)
	return num/denum # double_area / side_length == height


def groupLines(points,minCount,tolerance):
	random.shuffle(points)
	sample=points[:57]
	for (i,a) in enumerate(sample):
		for (j,b) in enumerate(sample):
			if j<=i: continue
			ab=Line(a,b)
			for c in points:
				if c is a or c is b: continue
				if point2lineDistance(a,b,c)<=tolerance:
					ab.points.add(c)
			if len(ab.points)>=minCount:
				yield ab


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)
	kernel=np.ones((3,3),np.uint8)
	maskB=cv.inRange(rect,centers[0]-unit,centers[0]+unit)
	maskB=cv.morphologyEx(maskB,cv.MORPH_OPEN,kernel,iterations=1)
	maskB=cv.erode(maskB,kernel,iterations=2)
	maskW=cv.inRange(rect,centers[1]-unit,centers[1]+unit)
	maskW=cv.erode(maskW,kernel,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)
	stoneLocs=filterStones(contours,stones,stoneDims)

	linesImg=cv.cvtColor(np.zeros((h,w),np.uint8),cv.COLOR_GRAY2BGR)
	cv.drawContours(linesImg,[c for (point,c) in stoneLocs],-1,(255,255,255),-1)
	for (p,c) in stoneLocs:
		cv.drawMarker(linesImg,(int(p.x),int(p.y)),(255,0,0),cv.MARKER_CROSS)

	lineDict=dict()
	minCount=min(max(math.sqrt(len(stoneLocs))-2,3),7)
	print("min count:",minCount)
	for line in groupLines([point for (point,contour) in stoneLocs],minCount,2):
		key=line.getSortedPoints()
		if key in lineDict: # we already have a line with the same incident points
			continue
		lineDict[line.getSortedPoints()]=line
		obsolete=set()
		for ab in lineDict.values():
			if ab is line: continue
			if line.points<ab.points: # == impossible
				del lineDict[line.getSortedPoints()]
				break
			if ab.points<line.points:
				obsolete.add(ab.getSortedPoints())
		for key in obsolete: del lineDict[key]

	for line in sorted(lineDict, key=len, reverse=True)[:16]:
		print(len(line),line)
		(xa,ya)=line[0]
		(xb,yb)=line[-1]
		cv.line(linesImg,(int(xa),int(ya)),((int(xb),int(yb))),(255,255,0),1)
	show(linesImg)