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,homogenize
from analyzer.grid import transformPoint

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 computeAngle(self,line):
		ab=self.a-self.b
		cd=line.a-line.b
		alpha=math.atan(ab.y/ab.x)
		gamma=math.atan(cd.y/cd.x)
		fi=max(alpha,gamma)-min(alpha,gamma)
		return min(fi,math.pi-fi)

	def intersect(self,line):
		p=self.toProjective()
		q=line.toProjective()
		return EPoint.fromProjective(np.cross(p,q))

	def toProjective(self):
		return homogenize(np.cross(self.a.toProjective(),self.b.toProjective()))

	def transform(self,matrix):
		a=EPoint.fromProjective(transformPoint(self.a.toProjective(),matrix))
		b=EPoint.fromProjective(transformPoint(self.b.toProjective(),matrix))
		if a is None or b is None: return None
		return Line(a,b)

	def __str__(self): return "({0},{1})".format(self.a,self.b)
	def __repr__(self): return "Line({0},{1})".format(repr(self.a),repr(self.b))


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


def computeRectiMatrix(p,q,r,s):
	# p || q, r || s
	vanish1=homogenize(np.cross(p.toProjective(),q.toProjective()))
	vanish2=homogenize(np.cross(r.toProjective(),s.toProjective()))
	horizon=homogenize(np.cross(vanish1,vanish2))

	return np.matrix([horizon,[0,1,0],[0,0,1]])


def scoreMatrix(matrix,p,r,lines):
	p_=p.transform(matrix)
	r_=r.transform(matrix)
	if p_ is None or r_ is None:
		return math.inf
	score=0
	for ab in lines:
		if ab is p or ab is r: continue
		ab_=ab.transform(matrix)
		if ab_ is None:
			score+=math.pi/2
			continue
		alpha=min(ab_.computeAngle(p_), ab_.computeAngle(r_))
		if alpha<=math.pi/30:
			score+=alpha
		else: score+=math.pi/2
	return score


def groupParallels(lines,tolerance,w):
	ctrl=False
	for (i,p) in enumerate(lines):
		for (j,q) in enumerate(lines[i+1:]):
			a=p.intersect(q)
			if a.y>0 and a.x<w: continue
			for r in lines[i+1+j+1:]:
				b=r.intersect(p) # !! ideal points
				c=r.intersect(q)
				ab=b-a
				ac=c-a
				if abs(ab.x*ac.y-ab.y*ac.x)/2<=tolerance: # area
					yield (p,q,r)
					ctrl=True
					break
			if ctrl: break
		if ctrl:
			ctrl=False
			continue


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=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)
	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()
	# for row in res: print(row)
	# linePack=[
	# 	Line(EPoint(174.457,166.127),EPoint(174.96,27.253)),
	# 	Line(EPoint(191.333,38.075),EPoint(191.485,71.227)),
	# 	Line(EPoint(210.117,167.092),EPoint(205.0,50.0)),
	# 	Line(EPoint(127.7,25.6),EPoint(120.172,179.405)),
	# 	Line(EPoint(127.809,58.481),EPoint(124.324,127.427)),
	# 	Line(EPoint(85.964,191.64),EPoint(97.68,14.477)),
	# 	Line(EPoint(56.447,124.662),EPoint(54.889,137.918))
	# ]
	# linePack=[
	# 	Line(EPoint(56.447,124.662),EPoint(139.695,128.104)),
	# 	Line(EPoint(288.267,74.433),EPoint(191.485,71.227)),
	# 	Line(EPoint(252.926,29.71),EPoint(174.96,27.253)),
	# 	Line(EPoint(274.412,120.07),EPoint(41.065,112.759)),
	# 	Line(EPoint(289.674,108.019),EPoint(26.17,100.538)),
	# 	Line(EPoint(240.702,107.818),EPoint(41.065,112.759)),
	# 	Line(EPoint(174.457,166.127),EPoint(88.192,164.5))
	# ]
	# for (i,p) in enumerate(linePack):
	# 	for q in linePack[i+1:]:
	# 		print(p.intersect(q))
	# show(linesImg)

	# parallels=[]
	# for (i,ab) in enumerate(lines):
	# 	for cd in lines[i+1:]:
	# 		if ab.computeAngle(cd)>math.pi/4:
	# 			continue
	# 		parallels.append((ab,cd))
	# print("parallel lines candidate pairs:",len(parallels))

	# parallels=list(groupParallels(lines,50,w))
	# print("parallel triples:",len(parallels))
	# for (p,q,r) in parallels:
	# 	print(p,q,r)
	# 	print(p.intersect(q))
	# 	print(p.intersect(r))
	# 	print(q.intersect(r))
	# 	print()
	# 	matsImg_=np.copy(matsImg)
	# 	for pi in (p,q,r):
	# 		cv.line(matsImg_,(int(pi.a.x),int(pi.a.y)),(int(pi.b.x),int(pi.b.y)),(0,255,255),1)
	# 	show(matsImg_)

	# p=Line(corners[0],corners[1])
	# q=Line(corners[2],corners[3])
	# r=Line(corners[1],corners[2])
	# s=Line(corners[3],corners[0])
	#
	# matrix=computeRectiMatrix(p,q,r,s)
	# print(matrix)
	# for (point,contour) in stoneLocs:
	# 	point_=EPoint.fromProjective(transformPoint(point.toProjective(),matrix))
	# 	# print(point,"->",point_)
	# 	cv.line(matsImg,(int(point.x),int(point.y)),(int(point_.x),int(point_.y)),(0,255,255),1)
	# points1=np.float32([[p.x-x1,p.y-y1] for p in corners])
	# points2=np.float32([[0,0],[0,400],[400,400],[400,0]])
	# print(points1)
	# print(points2)
	# mat=cv.getPerspectiveTransform(points1,points2)
	# print(mat)
	# show(rect)
	# warped=cv.warpPerspective(rect,matrix,(400,400))
	# show(warped)

	# mats=[]
	# for (i,(p,q)) in enumerate(parallels):
	# 	for (r,s) in parallels[i+1:]:
	# 		if p is r or p is s or q is r or q is s:
	# 			continue
	# 		matrix=computeRectiMatrix(p,q,r,s)
	# 		score=scoreMatrix(matrix,p,r,lines)
	# 		mats.append((score,p,q,r,s,matrix))
	# mats.sort(key=lambda x: x[0])
	# for (score,p,q,r,s,matrix) in mats[:4]:
	# 	print(score,p,q,r,s,matrix)
	# 	matsImg_=np.copy(matsImg)
	# 	for ab in (p,q,r,s):
	# 		((xa,ya),(xb,yb))=(ab.a,ab.b)
	# 		cv.line(matsImg_,(int(xa),int(ya)),(int(xb),int(yb)),(255,255,0),1)
	# 	show(matsImg_)
	# for (score,p,q,r,s,matrix) in mats[-4:]:
	# 	print(score,p,q,r,s,matrix)
	# 	matsImg_=np.copy(matsImg)
	# 	for ab in (p,q,r,s):
	# 		((xa,ya),(xb,yb))=(ab.a,ab.b)
	# 		cv.line(matsImg_,(int(xa),int(ya)),(int(xb),int(yb)),(0,0,255),1)
	# 	show(matsImg_)