import sys

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

import os

import math

import random

import logging as log

import cv2 as cv

import numpy as np

import scipy.cluster

import scipy.ndimage

import scipy.signal

from geometry import Line, point2lineDistance

from polar_hough import PolarHough

from annotations import DataFile,computeBoundingBox

from analyzer.epoint import EPoint

from analyzer.corners import Corners

random.seed(361)

log.basicConfig(level=log.DEBUG,format="%(message)s")

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)

	log.debug(colors)

	(centers,distortion)=scipy.cluster.vq.kmeans(arr,colors)

	log.debug("k-means centers: %s",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)

	log.debug("accepted: %s",len(res))

	log.debug("rejected: %s",len(contours)-len(res))

	return res

def groupLines(points,minCount,tolerance):

	random.shuffle(points)

	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

class BoardDetector:

	def __init__(self,annotationsPath):

		self._annotations=DataFile(annotationsPath)

		self._rectW=0

		self._rectH=0

		self._rect=None

	def __call__(self,img,filename):

		# approximately detect the board

		(h,w)=img.shape[:2]

		log.debug("image dimensions: %s x %s",w,h)

		show(img,filename)

		(x1,y1,x2,y2)=self._detectRough(img,filename)

		rect=img[y1:y2,x1:x2]

		self._rectW=x2-x1

		self._rectH=y2-y1

		self._rect=rect

		# quantize colors

		colors=self._sampleColors(rect)

		quantized=quantize(rect,colors)

		show(quantized,filename)

		# detect black and white stones

		stones=self._detectStones(quantized,colors)

	def _detectRough(self,img,filename):

		corners=self._annotations[filename][0]

		(x1,y1,x2,y2)=computeBoundingBox(corners)

		log.debug("bounding box: (%s,%s) - (%s,%s)",x1,y1,x2,y2)

		return (x1,y1,x2,y2)

	def _sampleColors(self,rect):

		(h,w)=rect.shape[:2]

		minirect=rect[h//4:3*h//4, w//4:3*w//4]

		return kmeans(minirect)

	def _detectStones(self,quantized,colors):

		(h,w)=quantized.shape[:2]

		mask=self._maskStones(quantized,colors)

		stoneDims=(w/19,h/19)

		log.debug("stone dims: %s - %s",tuple(x/2 for x in stoneDims),stoneDims)

		(contours,hierarchy)=cv.findContours(mask,cv.RETR_LIST,cv.CHAIN_APPROX_SIMPLE)

		stoneLocs=filterStones(contours,mask,stoneDims)

		return stoneLocs

	def _maskStones(self,quantized,colors):

		unit=np.array([1,1,1],dtype=np.uint8)

		kernel=np.ones((3,3),np.uint8)

		maskB=cv.inRange(quantized,colors[0]-unit,colors[0]+unit)

		maskB=cv.morphologyEx(maskB,cv.MORPH_OPEN,kernel,iterations=1)

		maskB=cv.erode(maskB,kernel,iterations=3)

		# distTransform = cv.distanceTransform(maskB,cv.DIST_L2,5)

		# maskB=cv.inRange(distTransform,6,10)

		maskW=cv.inRange(quantized,colors[1]-unit,colors[1]+unit)

		maskW=cv.morphologyEx(maskW,cv.MORPH_OPEN,kernel,iterations=1)

		maskW=cv.erode(maskW,kernel,iterations=2)

		stones=cv.bitwise_or(maskB,maskW)

		return stones

	def _constructLines(self,stoneLocs):

		lineDict=dict()

		log.debug("min count: %s",minCount)

			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]

		log.debug("valid lines: %s",len(lineDict))

		lines=sorted(lineDict.values(), key=lambda ab: len(ab.points), reverse=True)

		# visualize

		linesImg=cv.cvtColor(np.zeros((self._rectH,self._rectW),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)

		for line in lines:

			points=line.getSortedPoints()

			(xa,ya)=points[0]

			(xb,yb)=points[-1]

			cv.line(linesImg,(int(xa),int(ya)),(int(xb),int(yb)),(255,255,0),1)

		show(linesImg)

		return lines

		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<=self._rectW and 0<=point.y<=self._rectH: continue

				polarHough.put(point.toPolar(imgCenter))

		log.debug(vanish)

		return vanish

	def _computeTransformationMatrix(self,vanish,lines):

		(v1,v2)=vanish

		(p,r)=sorted(lines,key=lambda p: point2lineDistance(p.a,p.b,v1))[:2]

		(q,s)=sorted(lines,key=lambda p: point2lineDistance(p.a,p.b,v2))[:2]

		(a,b,c,d)=Corners([p.intersect(q),q.intersect(r),r.intersect(s),s.intersect(p)]) # canonize the abcd order

		a_=EPoint(b.x,min(a.y,d.y))

		b_=EPoint(b.x,max(b.y,c.y))

		c_=EPoint(c.x,max(b.y,c.y))

		d_=EPoint(c.x,min(a.y,d.y))

		abcd=[list(point) for point in (a,b,c,d)]

		abcd_=[list(point) for point in (a_,b_,c_,d_)]

		log.debug("abcd: %s ->",(a,b,c,d))

		log.debug("-> abcd_: %s",(a_,b_,c_,d_))

		matrix=cv.getPerspectiveTransform(np.float32(abcd),np.float32(abcd_))

		log.debug("transformation matrix: %s",matrix)

		rect=np.copy(self._rect)

		for point in (a,b,c,d):

			cv.drawMarker(rect,(int(point.x),int(point.y)),(0,255,255),cv.MARKER_TILTED_CROSS)

		show(rect)

		transformed=cv.warpPerspective(rect,matrix,(self._rectW,self._rectH))

		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 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 angleDiff(alpha,beta):

	diff=abs(alpha-beta)

	if diff>math.pi: diff=2*math.pi-diff

	return diff

import os

import sys

import numpy as np

import cv2 as cv

from annotations import DataFile,computeBoundingBox

def show(img,filename="x"):

	cv.imshow(filename,img)

	cv.waitKey(0)

	cv.destroyAllWindows()

def filterVert(edges):

	kernel = np.array([[1,0,1],[1,0,1],[1,0,1]],np.uint8)

	edges = cv.erode(edges,kernel)

	kernel=np.array([[0,1,0],[0,1,0],[0,1,0]],np.uint8)

	edges=cv.dilate(edges,kernel)

	return edges

def filterHor(edges):

	kernel = np.array([[1,1,1],[0,0,0],[1,1,1]],np.uint8)

	edges = cv.erode(edges,kernel)

	kernel=np.array([[0,0,0],[1,1,1],[0,0,0]],np.uint8)

	edges=cv.dilate(edges,kernel)

	return edges

def filterDiag(edges):

	kernel = np.array([[0,0,1],[1,0,0],[0,1,0]],np.uint8)

	edges1 = cv.erode(edges,kernel)

	kernel=np.array([[1,0,0],[0,1,0],[0,0,1]],np.uint8)

	edges1=cv.dilate(edges1,kernel)

	kernel = np.array([[0,1,0],[1,0,0],[0,0,1]],np.uint8)

	edges2 = cv.erode(edges,kernel)

	kernel=np.array([[0,0,1],[0,1,0],[1,0,0]],np.uint8)

	edges2=cv.dilate(edges2,kernel)

	return edges1+edges2

def houghLines(bwImg):

	colorImg=cv.cvtColor(bwImg,cv.COLOR_GRAY2BGR)

	lines = cv.HoughLinesP(bwImg,1,np.pi/180,10,minLineLength=10,maxLineGap=40)

	if lines is None: lines=[]

	for line in lines:

		x1,y1,x2,y2 = line[0]

		cv.line(colorImg,(x1,y1),(x2,y2),(0,255,0),1)

	show(colorImg)

if __name__=="__main__":

	i=sys.argv[1]

	annotations=DataFile("/home/laman/Projekty/python/oneEye/images/annotations.json.gz")

	filename="{0}.jpg".format(i)

	img=cv.imread(os.path.join("/home/laman/Projekty/python/oneEye/images/",filename))

	(x1,y1,x2,y2)=computeBoundingBox(annotations[filename][0])

	img=img[y1:y2, x1:x2, :]

	# blurred=cv.GaussianBlur(img,(5,5),0)

	# small=cv.resize(img,None,fx=0.5,fy=0.5,interpolation=cv.INTER_AREA)

	small=img

	clahe = cv.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))

	gray=cv.cvtColor(small,cv.COLOR_BGR2GRAY)

	# gray=clahe.apply(gray)

	show(gray)

	edges=cv.Canny(gray,70,130)

	show(edges)

	edges=filterHor(edges)+filterVert(edges)+filterDiag(edges)

	show(edges)

	# kernel = np.ones((2,2),np.uint8)

	# edges = cv.morphologyEx(edges, cv.MORPH_DILATE, kernel)

	# show(edges)

	# edges=cv.morphologyEx(edges,cv.MORPH_ERODE,kernel)

	# show(edges)

	colorEdges=cv.cvtColor(edges,cv.COLOR_GRAY2BGR)

import itertools

import math

import logging as log

import numpy as np

import scipy.signal

from geometry import angleDiff

class PolarHough:

	# http://www.cis.pku.edu.cn/vision/vpdetection/LiPYZ10isvc.pdf

	def __init__(self,anglePrecision,lengthPrecision):

		self._anglePrecision=anglePrecision

		self._lengthPrecision=lengthPrecision

		self._maxLength=4096

		n=math.ceil(2*math.pi/anglePrecision)

		self._acc=[[] for i in range(n)]

	def put(self,item):

		k=int(item[0]//self._anglePrecision)

		self._acc[k].append(item)

		vanishingPoints=[]

		angles=[alpha for (alpha,prominence) in angles]

		bins=self._mapAngles(angles)

		for (alpha,bin) in zip(angles,bins):

			(length,sampleCount)=self._extractLength([dist for (beta,dist) in bin])

			vanishingPoints.append((alpha,length))

		return vanishingPoints

		(peakKeys,info)=scipy.signal.find_peaks(lens,prominence=0)

		res=sorted(zip(info["prominences"],peakKeys),reverse=True)[:k]

		log.debug("(angle, prominence): %s ... %s",res,[alpha/self._anglePrecision for (alpha,_) in res])

		return res

	def _mapAngles(self,angles):

		res=[[] for alpha in angles]

		for (i,bin) in enumerate(self._acc):

			beta=i*self._anglePrecision

			key=min(zip(map(lambda alpha: angleDiff(alpha, beta), angles), itertools.count()))[1]

			res[key].extend(bin)

		return res

	def _extractLength(self,arr):

		acc=np.zeros(self._maxLength+1,dtype=np.int32)

		for dist in arr:

			dist=min(dist,self._maxLength)

			acc[int(dist//self._lengthPrecision)]+=1

		res=acc.argmax()

		log.debug("(length, count): %s",(res*self._lengthPrecision,acc[res]))

		return (res*self._lengthPrecision,acc[res])