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

import math
from datetime import datetime
import logging as log

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

from geometry import Line
from analyzer.epoint import EPoint

DEBUG=True


class BaseHough:
	def __init__(self,width,height):
		self._diagLen=int(np.sqrt(height**2+width**2))+1
		self._center=(width//2,height//2)
		self._acc=np.zeros((180,self._diagLen),dtype=np.int32)

	def update(self,x,y,weight=1):
		""":param x: number, 0 <= x < width
		:param y: number, 0 <= y < height"""
		for alphaDeg in range(0,180):
			d=self._computeDist(x,y,alphaDeg)+self._diagLen//2
			self._acc[(alphaDeg,d)]+=weight

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

		img=self._createImg()
		img=self._markPeaks(img,self._filterClose(peaks))
		self.show(img)

		res=[]
		for (alphaDeg,d) in peaks:
			alphaRad=alphaDeg*math.pi/180
			baseLine=Line(alphaRad,0)
			dd=baseLine.distanceTo(EPoint(*self._center)) # to shift d from the center to 0,0
			res.append(Line(alphaRad, dd+d-shift))
		log.debug("detected lines: %s",res)
		return res

	def _computeDist(self,x,y,alphaDeg):
		"""Compute the distance of a line with the provided alphaDeg declination and passing the (x,y) point to the image center.
		The returned distance might be negative (meaning the angle is in fact alpha+180)."""
		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 show(self,img=None):
		if not DEBUG: return
		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


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)

		h2=BaseHough(self._diagLen,180+90)
		for (alpha,d) in peaks:
			h2.update(d+shift,alpha)
			if alpha<90:
				h2.update(shift-d,alpha+180)
		lines=h2.extract(3)

		img=self._createImg()
		img=self._markPeaks(img,self._filterClose(allPeaks[:38]))

		for (i,line) in enumerate(lines):
			self.drawLine(img,line,i)
		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=round(curve(x))
			if y<0 or y>=2*h: continue
			if y<h:	img[y,x]=color
			else: img[y%h,-x]=color

	def drawLine(self,img,line,colorKey=0):
		colors=[[0,255,255],[255,0,255],[255,255,0]]
		color=colors[colorKey]
		(h,w)=img.shape[:2]
		(a,b,c)=line.toNormal()
		print("%",a,b,c)
		if b==0: return
		for x in range(1,w,3):
			y=round((-c-a*x)/b) + (0 if b>=0 else 180)
			if y<0 or y>=h: continue
			img[y,x]=color


def findPeaks(arr2d): # a naive implementation
	"""Scan 8-neighbourhood and for each peak or top plateau yield one point. For plateaus yield the """
	(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)