import itertools
import math

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):
		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])