import itertools
import math
import logging as log

import numpy as np
import scipy.signal

from geometry import angleDiff


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,trueVs):
		vanishingPoints=[]
		angles=self._extractAngles(count,tuple(v[0] for v in trueVs))
		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,trueAngles):
		lens=np.array([0]+list(map(len,self._acc))+[0])
		marked=np.copy(lens[1:-1])
		for alpha in trueAngles:
			key=int(alpha/self._anglePrecision)
			marked[key]=-marked[key]-1
		log.debug(marked)
		(peakKeys,info)=scipy.signal.find_peaks(lens,prominence=0)
		res=sorted(zip(info["prominences"],peakKeys),reverse=True)[:k]
		res=[((key-1)*self._anglePrecision,prominence) for (prominence,key) in res]
		log.debug("(angle, prominence): %s ... %s",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()
		log.debug("(length, count): %s",(res*self._lengthPrecision,acc[res]))
		return (res*self._lengthPrecision,acc[res])