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 annotations import DataFile,computeBoundingBox

from hough import show,prepareEdgeImg,HoughTransform

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)

	wood=[193,165,116]

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

	log.debug("k-means centers: %s",centers)

	(black,empty,white)=sorted(centers,key=sum)

	if np.linalg.norm(black)>np.linalg.norm(black-wood):

		black=None

	if np.linalg.norm(white-[255,255,255])>np.linalg.norm(white-wood):

		white=None

	log.debug("black, white: %s, %s",black,white)

	return (black,white,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))

	show(contourImg,"accepted and rejected stones")

	return res

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

		(black,white,colors)=self._sampleColors(rect)

		quantized=quantize(rect,colors)

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

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

		show(edges,"edges")

		quantized=quantized & (255-cv.cvtColor(edges,cv.COLOR_GRAY2BGR))

		show(quantized,"quantized, edges separated")

		# detect black and white stones

		stones=self._detectStones(quantized,black,white)

		# detect lines from edges and stones

		edgeImg=prepareEdgeImg(rect)

		hough=HoughTransform(edgeImg)

		stonesImg=np.zeros((self._rectH,self._rectW),np.uint8)

		for (point,c) in stones:

			cv.circle(stonesImg,(int(point.x),int(point.y)),2,255,-1)

		# cv.drawContours(stonesImg,[c for (point,c) in stones],-1,255,-1)

		show(stonesImg,"detected stones")

		# # detect vanishing points of the lines

		# imgCenter=EPoint(w//2-x1, h//2-y1)

		# (a,b,c,d)=(p-EPoint(x1,y1) for p in self._annotations[filename][0])

		# (p,q,r,s)=(Line(a,b),Line(b,c),Line(c,d),Line(d,a))

		# v1=p.intersect(r)

		# v2=q.intersect(s)

		# log.debug("true vanishing points: %s ~ %s, %s ~ %s",v1,v1.toPolar(imgCenter),v2,v2.toPolar(imgCenter))

		# vanish=self._detectVanishingPoints(lines,imgCenter,(v1.toPolar(imgCenter),v2.toPolar(imgCenter)))

		#

		# # rectify the image

		# matrix=self._computeTransformationMatrix(vanish,lines)

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

		#

		# # determine precise board edges

	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,black,white):

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

		mask=self._maskStones(quantized,black,white)

		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,black,white):

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

		if black is not None:

			maskB=cv.inRange(quantized,black-unit,black+unit)

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

			maskB=cv.inRange(distTransform,6,20)

			show(maskB,"black areas")

		else: maskB=np.zeros(quantized.shape[:2],dtype=np.uint8)

		if white is not None:

			maskW=cv.inRange(quantized,white-unit,white+unit)

			distTransform=cv.distanceTransform(maskW,cv.DIST_L2,5)

			maskW=cv.inRange(distTransform,6,20)

			show(maskW,"white areas")

		else: maskW=np.zeros(quantized.shape[:2],dtype=np.uint8)

		stones=cv.bitwise_or(maskB,maskW)

		show(stones,"black and white areas")

		return stones

	def _printLines(self,lines,allPoints,img):

		for (i,line) in enumerate(lines):

			img_=np.copy(img)

			points=list(line.getSortedPoints())

			(a,b)=max(((a,b) for a in points for b in points if a<b),key=lambda ab: ab[0].dist(ab[1]))

			(xa,ya)=a

			(xb,yb)=b

			points.sort(key=lambda p: a.dist(p))

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

			cv.imwrite("/tmp/{0}.png".format(i),img_)

			pointDists=",".join(str(round(p1.dist(p2),3)) for (p1,p2) in zip(points[:-1],points[1:]))

			log.debug("\t".join(map(str,[i,line,line.score(allPoints),pointDists])))

	def _detectVanishingPoints(self,lines,imgCenter,trueVs):

		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

				# log.debug("%s -> %s",point,point.toPolar(imgCenter))

				polarHough.put(point.toPolar(imgCenter))

		vanish=[EPoint.fromPolar(p,imgCenter) for p in polarHough.extract(2,trueVs)]

		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))

		show(transformed)

		return matrix

if __name__=="__main__":

	detector=BoardDetector(sys.argv[2])

	filepath=sys.argv[1]

	filename=os.path.basename(filepath)

	img=cv.imread(filepath)

	detector(img,filename)

import math

from analyzer.epoint import EPoint

class Line:

	def __init__(self,alpha,d):

		self._alpha=alpha

		self._d=d

		self._sin=math.sin(alpha)

		self._cos=math.cos(alpha)

	@staticmethod

	def fromNormal(a,b,c):

		"""ax + by + c = 0"""

		(a_,b_,c_)=(a/norm,b/norm,c/norm)

		alpha=math.acos(a_) if b_>=0 else 2*math.pi-math.acos(a_)

		return Line(alpha,-c_)

	@staticmethod

	def fromPoints(a,b):

		return Line.fromNormal(a.y-b.y, b.x-a.x, (b.y-a.y)*a.x+(a.x-b.x)*a.y)

	def toNormal(self):

		# https://en.wikipedia.org/wiki/Line_(mathematics)#In_normal_form

		"""ax + by + c = 0"""

		return (self._cos, self._sin, -self._d)

	def intersect(self,line):

		if self._alpha==line._alpha: return None

		(a,b,c)=self.toNormal()

		(d,e,f)=line.toNormal()

		x=(b*f-c*e)/(a*e-b*d)

		y=(c*d-a*f)/(a*e-b*d)

		return EPoint(x,y)

	def distanceTo(self,point):

		# https://en.wikipedia.org/wiki/Point-line_distance#Line_defined_by_an_equation

		(a,b,c)=self.toNormal()

		return abs(a*point.x+b*point.y+c) # a**2 + b**2 == 1 for Hesse normal form

	def __str__(self): return "({0},{1})".format(self._alpha,self._d)

	def __repr__(self): return "Line({0},{1})".format(repr(self._alpha),repr(self._d))

def angleDiff(alpha,beta):

	diff=abs(alpha-beta)

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

	return diff

import sys

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

import math

from datetime import datetime

import logging as log

import numpy as np

import scipy.optimize

import scipy.signal

import cv2 as cv

DEBUG=True

class LineBag:

	def __init__(self):

		self._lines=[]

	def put(self,score,alpha,beta,peaks):

		self._lines.append((score,alpha,beta,peaks))

	def pull(self,count):

		self._lines.sort(reverse=True)

		res=[]

		for (score,alpha,beta,peaks) in self._lines:

			if any(abs(alpha-gamma)<10 and abs(beta-delta)<10 for (_,gamma,delta,_) in res): continue

			if any((beta-delta)!=0 and (alpha-gamma)/(beta-delta)<0 for (_,gamma,delta,_) in res): continue

			res.append((score,alpha,beta,peaks))

			if len(res)>=count: break

		return res

class HoughTransform:

	def __init__(self,img):

		self._angleBandwidth=30 # degrees

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

		self._diagLen=int(np.sqrt(h**2+w**2))+1

		self._center=(w//2,h//2)

		self._acc=np.zeros((180,self._diagLen),dtype=np.int32)

		self.update(img)

	def extract(self):

		img=self._createImg()

		self.show(img)

		i=0

			log.debug("score: %s",score)

			log.debug("alpha, beta: %s, %s",alpha,beta)

			self._drawLine(img,alpha,beta,peaks,i)

			i+=1

		self.show(img)

	def update(self,img,weight=1):

		start=datetime.now().timestamp()

		for (r,row) in enumerate(img):

			for (c,pix) in enumerate(row):

				if pix==0: continue

				for alphaDeg in range(0,180):

					d=self._computeDist(c,r,alphaDeg)+self._diagLen//2

					self._acc[(alphaDeg,d)]+=weight

		log.debug("Hough updated in %s s",round(datetime.now().timestamp()-start,3))

	def _computeDist(self,x,y,alphaDeg):

		alphaRad=alphaDeg*math.pi/180

		(x0,y0)=self._center

		(dx,dy)=(x-x0,y-y0)

		d=dx*math.cos(alphaRad)+dy*math.sin(alphaRad)

		return round(d)

	def _detectLines(self):

		bag=LineBag()

		for alpha in range(0,180,2):

			for beta in range(max(alpha-60,0),alpha+60,2):

				accLine=[self._acc[key] for key in self._readLineKeys(alpha,beta)]

				(peaks,props)=scipy.signal.find_peaks(accLine,prominence=0)

				(prominences,peaks)=zip(*sorted(zip(props["prominences"],peaks),reverse=True)[:19])

				bag.put(sum(prominences),alpha,beta,peaks)

		return bag.pull(2)

	def _readLineKeys(self,alpha,beta):

		n=self._diagLen-1

		res=[]

		for i in range(n+1):

			k=round((alpha*(n-i)+beta*i)/n)

			if k<0 or k>=180:

				k=k%180

				i=n+1-i

			res.append((k,i))

		return res

	def show(self,img=None):

		if img is None: img=self._createImg()

		show(img,"Hough transform accumulator")

	def _createImg(self):

		maxVal=self._acc.max()

		arr=np.expand_dims(np.uint8(255*self._acc//maxVal),axis=2)

		img=np.concatenate((arr,arr,arr),axis=2)

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

		for x in range(0,w,4): # y axis

			img[h//2,x]=[255,255,255]

		for y in range(0,h,4):

			img[y,w//2]=[255,255,255]

		return img

	def _drawLine(self,img,alpha,beta,peaks,colorKey):

		colors=[[0,255,255],[255,0,255],[255,255,0]]

		color=colors[colorKey]

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

		keys=self._readLineKeys(alpha,beta)

		for (y,x) in keys:

			if x%3!=0: continue

			if y<0 or y>=h: continue

			img[y,x]=color

		for k in peaks:

			(y,x)=keys[k]

			cv.drawMarker(img,(x,y),color,cv.MARKER_TILTED_CROSS,8)

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)