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, point2lineDistance
from polar_hough import PolarHough
from annotations import DataFile,computeBoundingBox
from hough import show
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)
	colors=np.array([[0,0,0],[255,255,255],[193,165,116]],np.float)
	log.debug(colors)
	(centers,distortion)=scipy.cluster.vq.kmeans(arr,colors)
	log.debug("k-means centers: %s",centers)
	return 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)
	return res


def groupLines(points,minCount,tolerance):
	random.shuffle(points)
	sample=points[:57]
	for (i,a) in enumerate(sample):
		for (j,b) in enumerate(sample):
			if j<=i: continue
			ab=Line(a,b)
			for c in points:
				if c is a or c is b: continue
				if point2lineDistance(a,b,c)<=tolerance:
					ab.points.add(c)
			if len(ab.points)>=minCount:
				yield ab


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
		colors=self._sampleColors(rect)
		quantized=quantize(rect,colors)
		show(quantized,filename)

		# detect black and white stones
		stones=self._detectStones(quantized,colors)

		# detect lines passing through the stones
		lines=self._constructLines(stones)

		# detect vanishing points of the lines
		imgCenter=EPoint(w//2-x1, h//2-y1)
		vanish=self._detectVanishingPoints(lines,imgCenter)
		(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))

		# 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,colors):
		(h,w)=quantized.shape[:2]
		mask=self._maskStones(quantized,colors)
		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,colors):
		unit=np.array([1,1,1],dtype=np.uint8)
		kernel=np.ones((3,3),np.uint8)
		maskB=cv.inRange(quantized,colors[0]-unit,colors[0]+unit)
		maskB=cv.morphologyEx(maskB,cv.MORPH_OPEN,kernel,iterations=1)
		maskB=cv.erode(maskB,kernel,iterations=3)
		# distTransform = cv.distanceTransform(maskB,cv.DIST_L2,5)
		# maskB=cv.inRange(distTransform,6,10)
		show(maskB)
		maskW=cv.inRange(quantized,colors[1]-unit,colors[1]+unit)
		maskW=cv.morphologyEx(maskW,cv.MORPH_OPEN,kernel,iterations=1)
		maskW=cv.erode(maskW,kernel,iterations=2)
		show(maskW)
		stones=cv.bitwise_or(maskB,maskW)
		show(stones)
		return stones

	def _constructLines(self,stoneLocs):
		lineDict=dict()
		minCount=min(max(math.sqrt(len(stoneLocs))-4,3),7)
		log.debug("min count: %s",minCount)
		for line in groupLines([point for (point,contour) in stoneLocs],minCount,2):
			key=line.getSortedPoints()
			if key in lineDict: # we already have a line with the same incident points
				continue
			lineDict[line.getSortedPoints()]=line
			obsolete=set()
			for ab in lineDict.values():
				if ab is line: continue
				if line.points<ab.points: # == impossible
					del lineDict[key]
					break
				if ab.points<line.points:
					obsolete.add(ab.getSortedPoints())
			for key in obsolete: del lineDict[key]
		log.debug("valid lines: %s",len(lineDict))
		lines=sorted(lineDict.values(), key=lambda ab: len(ab.points), reverse=True)

		# visualize
		linesImg=cv.cvtColor(np.zeros((self._rectH,self._rectW),np.uint8),cv.COLOR_GRAY2BGR)
		cv.drawContours(linesImg,[c for (point,c) in stoneLocs],-1,(255,255,255),-1)
		for (p,c) in stoneLocs:
			cv.drawMarker(linesImg,(int(p.x),int(p.y)),(255,0,0),cv.MARKER_CROSS)
		for line in lines:
			points=line.getSortedPoints()
			(xa,ya)=points[0]
			(xb,yb)=points[-1]
			cv.line(linesImg,(int(xa),int(ya)),(int(xb),int(yb)),(255,255,0),1)
		show(linesImg)

		return lines

	def _detectVanishingPoints(self,lines,imgCenter):
		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)]
		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)