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 polar_hough import PolarHough
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)
	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,"accepted and rejected stones")
	return res


def groupLines(points,minCount,tolerance):
	random.shuffle(points)
	sample=points[:]
	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)
		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,colors)

		# 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)
		# cv.drawContours(stonesImg,[c for (point,c) in stones],-1,255,-1)
		show(stonesImg,"detected stones")
		hough.update(stonesImg,5)
		hough.extract()

		# # detect lines passing through the stones
		# lines=self._constructLines(stones)
		#
		# # detect vanishing points of the lines
		# imgCenter=EPoint(w//2-x1, h//2-y1)
		# (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))
		# vanish=self._detectVanishingPoints(lines,imgCenter,(v1.toPolar(imgCenter),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)
		maskB=cv.inRange(quantized,colors[0]-unit,colors[0]+unit)

		distTransform=cv.distanceTransform(maskB,cv.DIST_L2,5)
		maskB=cv.inRange(distTransform,6,20)
		show(maskB,"black areas")

		maskW=cv.inRange(quantized,colors[1]-unit,colors[1]+unit)
		distTransform=cv.distanceTransform(maskW,cv.DIST_L2,5)
		maskW=cv.inRange(distTransform,6,20)

		show(maskW,"white areas")
		stones=cv.bitwise_or(maskB,maskW)
		show(stones,"black and white areas")
		return stones

	def _constructLines(self,stoneLocs):
		lineDict=dict()
		# minCount=min(max(math.sqrt(len(stoneLocs))-4,3),7)
		minCount=6
		log.debug("min count: %s",minCount)
		points=[point for (point,contour) in stoneLocs]
		for line in groupLines(points,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)
		self._printLines(lines,points,linesImg)
		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 _printLines(self,lines,allPoints,img):
		for (i,line) in enumerate(lines):
			img_=np.copy(img)
			points=list(line.getSortedPoints())
			(a,b)=max(((a,b) for a in points for b in points if a<b),key=lambda ab: ab[0].dist(ab[1]))
			(xa,ya)=a
			(xb,yb)=b
			points.sort(key=lambda p: a.dist(p))
			cv.line(img_,(int(xa),int(ya)),(int(xb),int(yb)),(255,255,0),1)
			cv.imwrite("/tmp/{0}.png".format(i),img_)
			pointDists=",".join(str(round(p1.dist(p2),3)) for (p1,p2) in zip(points[:-1],points[1:]))
			log.debug("\t".join(map(str,[i,line,line.score(allPoints),pointDists])))

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