import sys
sys.path.append("../src")

import os
import math
import heapq

import cv2 as cv
import numpy as np
import scipy.cluster
import scipy.ndimage

from annotations import DataFile,computeBoundingBox
from hough import show
from analyzer.epoint import EPoint


class NeighboringPoint(EPoint):
	def __init__(self,x,y):
		super().__init__(x,y)
		self.neighbours=[]


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)
	print(colors)
	(centers,distortion)=scipy.cluster.vq.kmeans(arr,colors)
	print("k-means centers:",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.3 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)
	print("accepted:",len(res))
	print("rejected:",len(contours)-len(res))
	show(contourImg)
	return res


def computeDistances(points):
	n=len(points)
	res=np.zeros((n,n),dtype=np.float32)

	for (i,a) in enumerate(points):
		for (j,b) in enumerate(points):
			if j<i: continue
			elif j==i: res[i,j]=0
			else:
				res[i,j]=res[j,i]=a.dist(b)

	return res


def collectNeighbours(points):
	# we could do this in O(n log n)
	# doi:10.1007/BF02187718
	# https://link.springer.com/content/pdf/10.1007/BF02187718.pdf
	""":param points: [EPoint, ...]
	:return: [(k_i,dist_i) for i in range(4)]"""
	neighborhood=[NeighboringPoint(p.x,p.y) for p in points]
	dists=computeDistances(points)
	for (i,a) in enumerate(points):
		neighbours=heapq.nsmallest(4,enumerate(dists[i]),key=lambda x: x[1] if x[1]>0 else 1000)
		neighborhood[i].neighbours=[neighborhood[j] for (j,d) in neighbours]
	return neighborhood


def growLine(r,s):
	""":param r: NeighboringPoint
	:param s: NeighboringPoint
	:return: NeighboringPoint or None"""
	u=s-r
	alpha=math.atan2(u.y,u.x)
	t=None
	beta=alpha+math.pi # -alpha
	for t_ in s.neighbours:
		v=t_-s
		beta_=math.atan2(v.y,v.x) # beta_ in (-pi, pi]
		if abs(abs(alpha-beta_)-math.pi)>abs(abs(alpha-beta)-math.pi):
			beta=beta_
			t=t_
	if abs(alpha-beta)<math.pi/12: return t
	else: return None


if __name__=="__main__":
	filepath=sys.argv[1]
	annotations=DataFile(sys.argv[2])
	filename=os.path.basename(filepath)
	(x1,y1,x2,y2)=computeBoundingBox(annotations[filename][0])
	(w,h)=(x2-x1,y2-y1)
	img=cv.imread(filepath)
	(x3,x4,y3,y4)=(x1+w//4,x1+3*w//4,y1+h//4,y1+3*h//4)
	print("x3,x4,y3,y4:",x3,x4,y3,y4)
	rect=img[y3:y4,x3:x4,:]
	centers=kmeans(rect)
	print("x1,x2,y1,y2:",(x1,x2,y1,y2))
	img[y1:y2,x1:x2,:]=quantize(img[y1:y2,x1:x2,:],centers)
	print("image quantized")

	rect=img[y1:y2,x1:x2]
	unit=np.array([1,1,1],dtype=np.uint8)
	maskB=cv.inRange(rect,centers[0]-unit,centers[0]+unit)
	maskB=cv.dilate(maskB,np.ones((3,3),np.uint8),iterations=1)
	maskB=cv.erode(maskB,np.ones((3,3),np.uint8),iterations=3)
	maskW=cv.inRange(rect,centers[1]-unit,centers[1]+unit)
	maskW=cv.erode(maskW,np.ones((3,3),np.uint8),iterations=2)

	show(img,filename)
	show(maskB,filename)
	show(maskW,filename)
	stones=cv.bitwise_or(maskB,maskW)
	show(stones)

	stoneDims=(w/19,h/19)
	print("stone dims:",tuple(x/2 for x in stoneDims),"-",stoneDims)

	(contours,hierarchy)=cv.findContours(stones,cv.RETR_LIST,cv.CHAIN_APPROX_SIMPLE)
	stoneLocs=filterStones(contours,stones,stoneDims)
	neighborhood=collectNeighbours([point for (point,contour) in stoneLocs])

	linesImg=cv.cvtColor(np.zeros((h,w),np.uint8),cv.COLOR_GRAY2BGR)
	cv.drawContours(linesImg,[c for (point,c) in stoneLocs],-1,(255,255,255),-1)
	for p in neighborhood:
		cv.drawMarker(linesImg,(int(p.x),int(p.y)),(255,0,0),cv.MARKER_CROSS)

	for p in neighborhood:
		for q in p.neighbours:
			length=1
			while True:
				q_=growLine(p,q)
				if q_:
					q=q_
					length+=1
				else: break
			if length>1 or True:
				print(length,p,q)
			cv.line(linesImg,(int(p.x),int(p.y)),(int(q.x),int(q.y)),(255,255,0),min((length+1)//2,3))
	show(linesImg)
	cv.imwrite("/tmp/img.png",linesImg)