import sys

from polar_hough import PolarHough

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

import os
import math
import random

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

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

random.seed(361)


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 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)
	corners=annotations[filename][0]
	(x1,y1,x2,y2)=computeBoundingBox(corners)
	(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=4)
	maskW=cv.inRange(rect,centers[1]-unit,centers[1]+unit)
	maskW=cv.erode(maskW,kernel,iterations=3)

	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)
	matsImg=np.copy(linesImg)

	lineDict=dict()
	minCount=min(max(math.sqrt(len(stoneLocs))-4,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[key]
				break
			if ab.points<line.points:
				obsolete.add(ab.getSortedPoints())
		for key in obsolete: del lineDict[key]

	print("valid lines:",len(lineDict))
	lines=sorted(lineDict.values(), key=lambda ab: len(ab.points), reverse=True)
	res=[]
	for line in lines:
		v=line.b-line.a
		alpha=math.atan(v.y/v.x)
		res.append((round(math.pi/2-alpha if alpha>0 else math.pi/2+alpha,3),repr(line)))
		(xa,ya)=line.a
		(xb,yb)=line.b
		cv.line(linesImg,(int(xa),int(ya)),(int(xb),int(yb)),(255,255,0),1)
	res.sort()
	show(linesImg)

	imgSize=img.shape[:2]
	print("image size:",imgSize)
	imgCenter=EPoint(imgSize[1]/2,imgSize[0]/2)-EPoint(x1,y1)
	polarHough=PolarHough(math.pi/180,10)
	for (i,ab) in enumerate(lines):
		for cd in lines[i+1:]:
			point=ab.intersect(cd)
			if 0<=point.x<=w and 0<=point.y<=h: continue
			print(point,"->",point.toPolar(imgCenter))
			polarHough.put(point.toPolar(imgCenter))
	vanish=[EPoint.fromPolar(p,imgCenter) for p in polarHough.extract(2)]
	print(vanish)
	(a,b,c,d)=corners
	(p,q,r,s)=(Line(a,b),Line(b,c),Line(c,d),Line(d,a))
	v1=p.intersect(r)
	v2=q.intersect(s)
	print("true vanishing points:",v1,"~",v1.toPolar(imgCenter),"and",v2,"~",v2.toPolar(imgCenter))