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)
		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
		intersections=[]
		for p in lines[0]:
			for q in lines[1]:
				intersections.append(p.intersect(q))
		sack=DiagonalRansac(intersections,19)
		diagonals=sack.extract(10,2000)
		log.debug("diagonals candidates: %s",diagonals)
		for line in diagonals:
			self._drawLine(linesImg,line,[0,255,255])
		show(linesImg,"diagonals candidates")

	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

	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)