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

import numpy as np
import cv2 as cv

from annotations import DataFile,computeBoundingBox


class HoughTransform:
	def __init__(self,img):
		(h,w)=img.shape[:2]
		diagLen=np.sqrt(h**2+w**2)
		self._center=(w//2,h//2)
		self._acc=np.zeros((360,int(diagLen//2)+1),dtype=np.int32)

		self.update(img)

	def extract(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)
		log.debug(sorted(list(findPeaks(self._acc)),key=lambda rc: self._acc[rc],reverse=True)[:2*19])
		show(img,"Hough transform accumulator")

	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,360):
					d=self._computeDist(c,r,alphaDeg)
					if d>=0: 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 int(d)


def findPeaks(arr2d): # 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()