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

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

	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

		# 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

		# rectify the image

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

		show(maskB)

		maskW=cv.inRange(quantized,colors[1]-unit,colors[1]+unit)

		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

if __name__=="__main__":

	detector=BoardDetector(sys.argv[2])

	filepath=sys.argv[1]

	filename=os.path.basename(filepath)

	img=cv.imread(filepath)

	detector(img,filename)