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

import os
import math
from datetime import datetime
import logging as log

import numpy as np
import scipy.optimize
import cv2 as cv

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


class HoughTransform:
	def __init__(self,img):
		self._angleBandwidth=30 # degrees

		(h,w)=img.shape[:2]
		self._diagLen=int(np.sqrt(h**2+w**2))+1
		self._center=(w//2,h//2)
		self._acc=np.zeros((180,self._diagLen),dtype=np.int32)

		self.update(img)

	def extract(self):
		shift=self._diagLen//2
		allPeaks=sorted(list(findPeaks(self._acc)),key=lambda rc: self._acc[rc],reverse=True)
		peaks=allPeaks[:38]
		peaks=[(alpha,d-shift) for (alpha,d) in peaks]
		peaks=self._filterClose(peaks)
		peaks.sort(key=lambda rc: rc[0])
		log.debug("detected peaks: %s",peaks)
		(alpha,beta)=self._detectDominantAngles(peaks)

		img=self._createImg()
		img=self._markPeaks(img,self._filterClose(allPeaks[:38]))
		doublePeaks=peaks+[(alpha+180,-d) for (alpha,d) in peaks]
		params=self._computeGridParams([(gamma,d+shift) for (gamma,d) in doublePeaks if alpha<=gamma<=alpha+self._angleBandwidth])
		self._drawGridCurve(img,params,0)
		params=self._computeGridParams([(gamma,d+shift) for (gamma,d) in doublePeaks if beta<=gamma<=beta+self._angleBandwidth])
		self._drawGridCurve(img,params,1)
		self.show(img)

	def update(self,img,weight=1):
		start=datetime.now().timestamp()
		for (r,row) in enumerate(img):
			for (c,pix) in enumerate(row):
				if pix==0: continue
				for alphaDeg in range(0,180):
					d=self._computeDist(c,r,alphaDeg)+self._diagLen//2
					self._acc[(alphaDeg,d)]+=weight
		log.debug("Hough updated in %s s",round(datetime.now().timestamp()-start,3))

	def _computeDist(self,x,y,alphaDeg):
		alphaRad=alphaDeg*math.pi/180
		(x0,y0)=self._center
		(dx,dy)=(x-x0,y-y0)
		d=dx*math.cos(alphaRad)+dy*math.sin(alphaRad)
		return round(d)

	def _filterClose(self,peaks): # a naive implementation
		"""Discard points with Euclidean distance on the original image lower than 10.
		From such pairs keep only the one with a higher value in the accumulator.
		This can delete a series of points. If a-b and b-c are close and a>b>c, only a is kept."""
		minDist=13
		center=EPoint(*self._center)
		res=[]
		for (alphaDeg,d) in peaks:
			alphaRad=alphaDeg*math.pi/180
			point=EPoint.fromPolar((alphaRad,d),center)
			ctrl=True
			for (betaDeg,e) in peaks:
				betaRad=betaDeg*math.pi/180
				point_=EPoint.fromPolar((betaRad,e),center)
				if point.dist(point_)<minDist and self._acc[(alphaDeg,d)]<self._acc[(betaDeg,e)]:
					ctrl=False
			if ctrl: res.append((alphaDeg,d))
		return res

	def _detectDominantAngles(self,peaks):
		angles=[alpha for (alpha,d) in peaks]
		n=len(angles)
		bandwidth=self._angleBandwidth
		k1=0
		k2=1
		histogram=[]
		while k1<n:
			while (k2<n and angles[k1]+bandwidth>angles[k2]) or (k2>=n and angles[k1]+bandwidth>angles[k2%n]+180):
				k2+=1
			histogram.append((angles[k1],k2-k1))
			k1+=1
		log.debug("angles histogram: %s",histogram)
		dominantAngles=sorted(histogram,key=lambda xy: xy[1],reverse=True)
		alpha=dominantAngles[0]
		dominantAngles=[beta for beta in dominantAngles if 180-bandwidth>abs(alpha[0]-beta[0])>bandwidth]
		beta=dominantAngles[0]
		log.debug("dominant angles: %s, %s",alpha,beta)
		return (alpha[0],beta[0])

	def _computeGridParams(self,lines):
		log.debug("computing grid parameters for: %s",lines)
		angles=[alpha for (alpha,d) in lines]
		dists=[d for (alpha,d) in lines]
		curve=lambda x,a,b,c,d: a*x**3+b*x**2+c*x+d
		(params,cov)=scipy.optimize.curve_fit(curve,dists,angles)
		log.debug("result: %s",params)
		return params

	def show(self,img=None):
		if img is None: img=self._createImg()

		show(img,"Hough transform accumulator")

	def _createImg(self):
		maxVal=self._acc.max()
		arr=np.expand_dims(np.uint8(255*self._acc//maxVal),axis=2)
		img=np.concatenate((arr,arr,arr),axis=2)

		(h,w)=img.shape[:2]

		for x in range(0,w,4): # y axis
			img[h//2,x]=[255,255,255]
		for y in range(0,h,4):
			img[y,w//2]=[255,255,255]

		return img

	def _markPeaks(self,img,peaks):
		colors=[[255,0,0],[255,255,0],[0,255,0],[0,255,255],[0,0,255]]
		for (i,(alpha,d)) in enumerate(peaks[:38]):
			cv.drawMarker(img,(d,alpha),colors[i//9],cv.MARKER_TILTED_CROSS)
		return img

	def _drawGridCurve(self,img,params,colorKey=0):
		colors=[[0,255,255],[255,0,255],[255,255,0]]
		color=colors[colorKey]
		(a,b,c,d)=params
		(h,w)=img.shape[:2]
		curve=lambda x: a*x**3+b*x**2+c*x+d
		for x in range(0,w,3):
			y=int(curve(x))
			if y<0 or y>=2*h: continue
			if y<h:	img[y,x]=color
			else: img[y%h,-x]=color


def findPeaks(arr2d): # a naive implementation
	(h,w)=arr2d.shape
	neighbours=[(-1,-1),(-1,0),(-1,1),(0,-1),(0,1),(1,-1),(1,0),(1,1)]
	for r in range(h):
		for c in range(w):
			if all(r+dr<0 or r+dr>=h or c+dc<0 or c+dc>=w or arr2d[r,c]>arr2d[r+dr,c+dc] or (i<4 and arr2d[r,c]>=arr2d[r+dr,c+dc]) for (i,(dr,dc)) in enumerate(neighbours)):
				yield (r,c)


def show(img,filename="x"):
	cv.imshow(filename,img)
	cv.waitKey(0)
	cv.destroyAllWindows()


def filterVert(edges):
	kernel = np.array([[1,0,1],[1,0,1],[1,0,1]],np.uint8)
	edges = cv.erode(edges,kernel)
	kernel=np.array([[0,1,0],[0,1,0],[0,1,0]],np.uint8)
	edges=cv.dilate(edges,kernel)
	return edges

def filterHor(edges):
	kernel = np.array([[1,1,1],[0,0,0],[1,1,1]],np.uint8)
	edges = cv.erode(edges,kernel)
	kernel=np.array([[0,0,0],[1,1,1],[0,0,0]],np.uint8)
	edges=cv.dilate(edges,kernel)
	return edges

def filterDiag(edges):
	kernel = np.array([[0,0,1],[1,0,0],[0,1,0]],np.uint8)
	edges1 = cv.erode(edges,kernel)
	kernel=np.array([[1,0,0],[0,1,0],[0,0,1]],np.uint8)
	edges1=cv.dilate(edges1,kernel)

	kernel = np.array([[0,1,0],[1,0,0],[0,0,1]],np.uint8)
	edges2 = cv.erode(edges,kernel)
	kernel=np.array([[0,0,1],[0,1,0],[1,0,0]],np.uint8)
	edges2=cv.dilate(edges2,kernel)

	return edges1+edges2

def prepareEdgeImg(img):
	gray=cv.cvtColor(img,cv.COLOR_BGR2GRAY)
	show(gray,"greyscale image")
	edges=cv.Canny(gray,70,130)
	show(edges,"Canny edge detector")
	edges=filterHor(edges)+filterVert(edges)+filterDiag(edges)
	show(edges,"kernel filtered edges")
	return edges

def houghLines(bwImg):
	colorImg=cv.cvtColor(bwImg,cv.COLOR_GRAY2BGR)
	lines = cv.HoughLinesP(bwImg,1,np.pi/180,10,minLineLength=10,maxLineGap=40)
	if lines is None: lines=[]
	for line in lines:
		x1,y1,x2,y2 = line[0]
		cv.line(colorImg,(x1,y1),(x2,y2),(0,255,0),1)

	show(colorImg)


if __name__=="__main__":
	i=sys.argv[1]
	annotations=DataFile("/home/laman/Projekty/python/oneEye/images/annotations.json.gz")
	filename="{0}.jpg".format(i)
	img=cv.imread(os.path.join("/home/laman/Projekty/python/oneEye/images/",filename))
	(x1,y1,x2,y2)=computeBoundingBox(annotations[filename][0])
	img=img[y1:y2, x1:x2, :]
	# blurred=cv.GaussianBlur(img,(5,5),0)
	# small=cv.resize(img,None,fx=0.5,fy=0.5,interpolation=cv.INTER_AREA)
	small=img
	clahe = cv.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
	gray=cv.cvtColor(small,cv.COLOR_BGR2GRAY)
	# gray=clahe.apply(gray)
	show(gray)
	edges=cv.Canny(gray,70,130)
	show(edges)
	edges=filterHor(edges)+filterVert(edges)+filterDiag(edges)
	show(edges)


	# kernel = np.ones((2,2),np.uint8)
	# edges = cv.morphologyEx(edges, cv.MORPH_DILATE, kernel)
	# show(edges)
	# edges=cv.morphologyEx(edges,cv.MORPH_ERODE,kernel)
	# show(edges)
	colorEdges=cv.cvtColor(edges,cv.COLOR_GRAY2BGR)

	# houghLines(edges)
	h=HoughTransform(edges)
	h.extract()