import random

from geometry import Line


class LineBag:
	def __init__(self,capacity):
		self._capacity=capacity
		self._bag=[]

	def put(self,item):
		# early abort
		if self._bag and item>self._bag[-1]: return

		(v,p)=item
		# mostly filter duplicates
		for (i,(vi,pi)) in enumerate(self._bag):
			if abs(p.alpha-pi.alpha)<0.02 and abs(p.d-pi.d)<5:
				if v<vi: self._bag[i]=item
				return

		# add new and remove the last
		self._bag.append(item)
		self._bag.sort()
		if len(self._bag)>self._capacity:
			self._bag.pop()

	def pull(self,count=0):
		return self._bag[:count] if count else self._bag[:]


class Ransac:
	def __init__(self,points):
		self._points=points

	def extract(self,count,iterations):
		bag=LineBag(count)

		for i in range(iterations):
			line=self._generateLine()
			score=self._scoreLine(line)
			bag.put((score+random.random()/1000,line))

		return [line for (score,line) in bag.pull()]

	def _generateLine(self):
		(a,b)=random.sample(self._points,2)
		return Line.fromPoints(a,b)

	def _scoreLine(self,line):
		return sum(min(line.distanceTo(p),5) for p in self._points)


class DiagonalRansac(Ransac):
	def __init__(self,points,sideLen):
		super().__init__(points)
		self._sideLen=sideLen

	def _generateLine(self):
		n=len(self._points)
		m=self._sideLen
		ka=random.randrange(0,n)
		kb=ka

		while ka//m==kb//m or ka%m==kb%m:
			kb=random.randrange(0,n)
		a=self._points[ka]
		b=self._points[kb]

		return Line.fromPoints(a,b)