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

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)

		show(stonesImg,"detected stones")

		hough.update(stonesImg,10)

		lines=hough.extract()

		linesImg=np.copy(rect)

			self._drawLine(linesImg,line)

		show(linesImg,"detected lines")

		# # rectify the image

	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

		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 matrix

	def _drawLine(self,img,line):

		(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)),[0,255,0])

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

			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)

			keys=self._readLineKeys(alpha,beta)

			for k in peaks:

				(alphaDeg,d)=keys[k]

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

			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