import sys

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

import os

import random

import itertools

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

from ransac import DiagonalRansac

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)

		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)

		show(stonesImg,"detected stones")

		linesImg=np.copy(rect)

		for line in itertools.chain(*lines):

			self._drawLine(linesImg,line)

		show(linesImg,"detected lines")

		# # rectify the image

		matrix=self._computeTransformationMatrix(lines[0][0],lines[0][-1],lines[1][0],lines[1][-1])

		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 _computeTransformationMatrix(self,p,q,r,s): # p || q, r || s

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

		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,"rectified image")

		return matrix

	def _drawLine(self,img,line,color=None):

		if not color: color=[0,255,0]

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

		corners=[EPoint(0,0),EPoint(w,0),EPoint(0,h),EPoint(w,h)] # NW NE SW SE

		borders=[

			[Line.fromPoints(corners[0],corners[1]), Line.fromPoints(corners[2],corners[3])], # N S

			[Line.fromPoints(corners[0],corners[2]), Line.fromPoints(corners[1],corners[3])] # W E

		]

		(a,b)=(line.intersect(borders[0][0]), line.intersect(borders[0][1]))

		log.debug("%s %s",line,(a,b))

		if not a or not b:

			(a,b)=(line.intersect(borders[1][0]), line.intersect(borders[1][1]))

			log.debug("* %s %s",line,(a,b))

		if any(abs(x)>10**5 for x in [*a,*b]):

			log.debug("ignored")

			return

		cv.line(img,(int(a.x),int(a.y)),(int(b.x),int(b.y)),color)

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

from geometry import EPoint,Line

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

			# avoid intersecting lines

			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:

	"""Find line sequences with Hough transform.

	Uses usual image coordinates on input and output, with [0,0] in the upper left corner and [height-1,width-1] in the lower right.

	However, internally it uses the usual cartesian coordinates, centered at the image center. [-w/2,-h/2] in the upper left and [w/2,h/2] in the lower right."""

	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)

		lines=self._detectLines()

		res=[]

		i=0

		for (score,alpha,beta,peaks) in lines:

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

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

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

			res.append([])

			keys=self._readLineKeys(alpha,beta)

			for k in peaks:

				(alphaDeg,d)=keys[k]

				line=Line(alphaDeg*math.pi/180,d-self._diagLen//2)

				res[-1].append(self._transformOutput(line))

			i+=1

		self.show(img)

		return res

	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 show(self,img=None):

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

		show(img,"Hough transform accumulator")

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

		alphaRad=alphaDeg*math.pi/180

		(x0,y0)=self._center

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

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

		return round(d)

	def _detectLines(self):

		bag=LineBag()

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

			for beta in range(max(alpha-60,0),min(alpha+60,180+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>=180:

				k=k%180

				i=n-i

			res.append((k,i))

		return res

	def _transformOutput(self,line):

		(x,y)=self._center

		basis=EPoint(-x,y)

		shiftedLine=line.shiftBasis(basis)

		reflectedLine=Line(math.pi*2-shiftedLine.alpha,shiftedLine.d)

		log.debug("%s -> %s",line,reflectedLine)

		return reflectedLine

	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)

	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 prepareEdgeImg(img):

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

	show(gray,"greyscale image")

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

	show(edges,"Canny edge detector")

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

	show(edges,"kernel filtered edges")

	return edges

import logging

from .epoint import EPoint

log=logging.getLogger(__name__)

class Corners:

	def __init__(self,cornerList=[]):

		self._corners=cornerList[:]

		self._is_canon=False

		self._canonizeOrder()

	## Adds a new corner if there are less than four, replaces the closest otherwise.

	def add(self,x,y):

		a=EPoint(x,y)

		# for i,c in enumerate(self.corners): # move an improperly placed point

			# if a.dist(c)<20:

				# self.corners[i]=a

				# return

		if len(self._corners)<4: # add a new corner

			self._corners.append(a)

		if len(self._corners)<4:

			return

		index,minDist=0,float('inf') # replace the corner closest to the clicked point

		for i,c in enumerate(self._corners):

			if a.dist(c)<minDist:

				index,minDist=i,a.dist(c)

		self._corners[index]=a

		self._canonizeOrder()

	## Order the corners (0,1,2,3) so they make a quadrangle with vertices KLMN in counter-clockwise order, K being in the upper left.

	#

	#  For four points ABCD, there are 24 possible permutations corresponding to the desired KLMN.

	#  When we relax the condition of K being the upper left one, we get six groups of four equivalent permutations. KLMN ~ LMNK ~ MNKL ~ NKLM.

	#

	#  We determine which of the points' triplets are oriented clockwise and which counter-clockwise (minus/plus in the table below)

	#  and swap them so that all triangles turn counter-clockwise.

	#

	#  xxxx -> KLMN | ABC | ABD | ACD | BCD | index | swap

	#  ------------ | :-: | :-: | :-: | :-: | ----: | ----

	#  A BCD        |  +  |  +  |  +  |  +  |    15 | 0

	#  A BDC        |  +  |  +  |  -  |  -  |    12 | CD

	#  A CBD        |  -  |  +  |  +  |  -  |     6 | BC

	#  A CDB        |  -  |  -  |  +  |  +  |     3 | AB

	#  A DBC        |  +  |  -  |  -  |  +  |     9 | AD

	#  A DCB        |  -  |  -  |  -  |  -  |     0 | BD

	#

	#  For every non-degenerate quadrangle, there must be 1-3 edges going right-left (from a higher to a lower x coordinate).

	#  From these pick the one with the lowest slope (dy/dx) and declare its ending point the upper left corner. For the same slope pick the one further left.

	#

	#  @return True for a convex quadrangle, False for concave and degenerate cases.

	def _canonizeOrder(self):

		def false():

			self._is_canon=False

			return False

		if len(self._corners)!=4: return false()

		a,b,c,d=self._corners

		abc=doubleTriangleArea(a,b,c)

		abd=doubleTriangleArea(a,b,d)

		acd=doubleTriangleArea(a,c,d)

		bcd=doubleTriangleArea(b,c,d)

		if any(x==0 for x in (abc,abd,acd,bcd)): return false() # collinear degenerate

		swaps=[(1,3),(0,1),(1,2),(0,3),(2,3),(0,0)]

		index=(8 if abc>0 else 0)|(4 if abd>0 else 0)|(2 if acd>0 else 0)|(1 if bcd>0 else 0)

		if index%3!=0: return false() # concave degenerate

		swap=swaps[index//3]

		self._corners[swap[0]], self._corners[swap[1]] = self._corners[swap[1]], self._corners[swap[0]] # counter-clockwise order

		kIndex=None

		lowestSlope=float("inf")

		for i,corner in enumerate(self._corners): # find the NK edge: going right-left with the lowest slope, secondarily the one going down

			ii=(i+1)%4

			slope=abs(getSlope(corner, self._corners[ii]))

			if corner.x>self._corners[ii].x and (slope < lowestSlope or (slope == lowestSlope and corner.y < self._corners[ii].y)):

				kIndex=ii

				lowestSlope=slope

		self._corners= self._corners[kIndex:] + self._corners[:kIndex] # rotate the upper left corner to the first place

		self._is_canon=True # success

		return True

	def scale(self,scale):

		self._corners=[c * scale for c in self._corners]

	def __iter__(self):

		return iter(self._corners)

	def __len__(self):

		return len(self._corners)

	def is_canon(self):

		return self._is_canon

## Computes twice the area of the triangle formed by points a,b,c.

#

#  @return positive value for points oriented counter-clockwise, negative for clockwise, zero for degenerate cases.

def doubleTriangleArea(a,b,c):

	return (a.x-b.x)*(c.y-a.y)-(c.x-a.x)*(a.y-b.y)

def getSlope(a,b):

	if(b.x==a.x): return float("inf")

	return (b.y-a.y)/(b.x-a.x)