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

		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

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

		return matrix

		(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

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)

			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