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

import os
import math
import random
import itertools

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

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

random.seed(361)


class Line():
	def __init__(self,a,b):
		self.a=a
		self.b=b
		self.points={a,b}

	def getSortedPoints(self):
		return tuple(sorted(self.points))

	def computeAngle(self,line):
		ab=self.a-self.b
		cd=line.a-line.b
		alpha=math.atan(ab.y/ab.x)
		gamma=math.atan(cd.y/cd.x)
		fi=max(alpha,gamma)-min(alpha,gamma)
		return min(fi,math.pi-fi)

	def intersect(self,line):
		p=self.toProjective()
		q=line.toProjective()
		return EPoint.fromProjective(np.cross(p,q))

	def toProjective(self):
		return homogenize(np.cross(self.a.toProjective(),self.b.toProjective()))

	def transform(self,matrix):
		a=EPoint.fromProjective(transformPoint(self.a.toProjective(),matrix))
		b=EPoint.fromProjective(transformPoint(self.b.toProjective(),matrix))
		if a is None or b is None: return None
		return Line(a,b)

	def __str__(self): return "({0},{1})".format(self.a,self.b)
	def __repr__(self): return "Line({0},{1})".format(repr(self.a),repr(self.b))


class PolarHough:
	# http://www.cis.pku.edu.cn/vision/vpdetection/LiPYZ10isvc.pdf
	def __init__(self,anglePrecision,lengthPrecision):
		self._anglePrecision=anglePrecision
		self._lengthPrecision=lengthPrecision
		self._maxLength=4096
		n=math.ceil(2*math.pi/anglePrecision)
		self._acc=[[] for i in range(n)]

	def put(self,item):
		k=int(item[0]//self._anglePrecision)
		self._acc[k].append(item)

	def extract(self,count):
		vanishingPoints=[]
		angles=self._extractAngles(count)
		angles=[alpha for (alpha,prominence) in angles]
		bins=self._mapAngles(angles)
		for (alpha,bin) in zip(angles,bins):
			(length,sampleCount)=self._extractLength([dist for (beta,dist) in bin])
			vanishingPoints.append((alpha,length))
		return vanishingPoints

	def _extractAngles(self,k):
		lens=np.array(list(map(len,self._acc)))
		print(lens)
		(peakKeys,info)=scipy.signal.find_peaks(lens,prominence=0)
		res=sorted(zip(info["prominences"],peakKeys),reverse=True)[:k]
		res=[(key*self._anglePrecision,prominence) for (prominence,key) in res]
		print("(angle, prominence):",res,"...",[alpha/self._anglePrecision for (alpha,_) in res])
		return res

	def _mapAngles(self,angles):
		res=[[] for alpha in angles]
		for (i,bin) in enumerate(self._acc):
			beta=i*self._anglePrecision
			key=min(zip(map(lambda alpha: angleDiff(alpha,beta), angles), itertools.count()))[1]
			res[key].extend(bin)
		return res

	def _extractLength(self,arr):
		acc=np.zeros(self._maxLength+1,dtype=np.int32)
		for dist in arr:
			dist=min(dist,self._maxLength)
			acc[int(dist//self._lengthPrecision)]+=1
		res=acc.argmax()
		print("(length, count):",(res*self._lengthPrecision,acc[res]))
		return (res*self._lengthPrecision,acc[res])


def angleDiff(alpha,beta):
	diff=abs(alpha-beta)
	if diff>math.pi: diff=2*math.pi-diff
	return diff


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.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)
	print("accepted:",len(res))
	print("rejected:",len(contours)-len(res))
	show(contourImg)
	return res


def point2lineDistance(a,b,p):
	# https://en.wikipedia.org/wiki/Point-line_distance#Line_defined_by_two_points
	ab=b-a
	num=abs(ab.y*p.x - ab.x*p.y + b.x*a.y - a.x*b.y)
	denum=math.sqrt(ab.y**2+ab.x**2)
	return num/denum # double_area / side_length == height


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


def computeRectiMatrix(p,q,r,s):
	# p || q, r || s
	vanish1=homogenize(np.cross(p.toProjective(),q.toProjective()))
	vanish2=homogenize(np.cross(r.toProjective(),s.toProjective()))
	horizon=homogenize(np.cross(vanish1,vanish2))

	return np.matrix([horizon,[0,1,0],[0,0,1]])


def scoreMatrix(matrix,p,r,lines):
	p_=p.transform(matrix)
	r_=r.transform(matrix)
	if p_ is None or r_ is None:
		return math.inf
	score=0
	for ab in lines:
		if ab is p or ab is r: continue
		ab_=ab.transform(matrix)
		if ab_ is None:
			score+=math.pi/2
			continue
		alpha=min(ab_.computeAngle(p_), ab_.computeAngle(r_))
		if alpha<=math.pi/30:
			score+=alpha
		else: score+=math.pi/2
	return score


def groupParallels(lines,tolerance,w):
	ctrl=False
	for (i,p) in enumerate(lines):
		for (j,q) in enumerate(lines[i+1:]):
			a=p.intersect(q)
			if a.y>0 and a.x<w: continue
			for r in lines[i+1+j+1:]:
				b=r.intersect(p) # !! ideal points
				c=r.intersect(q)
				ab=b-a
				ac=c-a
				if abs(ab.x*ac.y-ab.y*ac.x)/2<=tolerance: # area
					yield (p,q,r)
					ctrl=True
					break
			if ctrl: break
		if ctrl:
			ctrl=False
			continue


if __name__=="__main__":
	filepath=sys.argv[1]
	annotations=DataFile(sys.argv[2])
	filename=os.path.basename(filepath)
	corners=annotations[filename][0]
	(x1,y1,x2,y2)=computeBoundingBox(corners)
	(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)
	kernel=np.ones((3,3),np.uint8)
	maskB=cv.inRange(rect,centers[0]-unit,centers[0]+unit)
	maskB=cv.morphologyEx(maskB,cv.MORPH_OPEN,kernel,iterations=1)
	maskB=cv.erode(maskB,kernel,iterations=4)
	maskW=cv.inRange(rect,centers[1]-unit,centers[1]+unit)
	maskW=cv.erode(maskW,kernel,iterations=3)

	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)

	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,c) in stoneLocs:
		cv.drawMarker(linesImg,(int(p.x),int(p.y)),(255,0,0),cv.MARKER_CROSS)
	matsImg=np.copy(linesImg)

	lineDict=dict()
	minCount=min(max(math.sqrt(len(stoneLocs))-4,3),7)
	print("min count:",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]

	print("valid lines:",len(lineDict))
	lines=sorted(lineDict.values(), key=lambda ab: len(ab.points), reverse=True)
	res=[]
	for line in lines:
		v=line.b-line.a
		alpha=math.atan(v.y/v.x)
		res.append((round(math.pi/2-alpha if alpha>0 else math.pi/2+alpha,3),repr(line)))
		(xa,ya)=line.a
		(xb,yb)=line.b
		cv.line(linesImg,(int(xa),int(ya)),(int(xb),int(yb)),(255,255,0),1)
	res.sort()
	show(linesImg)

	imgSize=img.shape[:2]
	print("image size:",imgSize)
	imgCenter=EPoint(imgSize[1]/2,imgSize[0]/2)-EPoint(x1,y1)
	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<=w and 0<=point.y<=h: continue
			print(point,"->",point.toPolar(imgCenter))
			polarHough.put(point.toPolar(imgCenter))
	vanish=[EPoint.fromPolar(p,imgCenter) for p in polarHough.extract(2)]
	print(vanish)
	(a,b,c,d)=corners
	(p,q,r,s)=(Line(a,b),Line(b,c),Line(c,d),Line(d,a))
	v1=p.intersect(r)
	v2=q.intersect(s)
	print("true vanishing points:",v1,"~",v1.toPolar(imgCenter),"and",v2,"~",v2.toPolar(imgCenter))