import pygame
from OpenGL.GL import *

def n2ceil(n):
	"Find the next power of two greater than or equal to n"
	assert n > 0
	x = 1
	while x <= 1024:
		if x >= n:
			return x
		x <<= 1
	# FAIL
	return 1024


def splitrect(rect, point):
	"Split a rect into four rects with point in the center"
	assert rect.collidepoint(point)
	return [
		pygame.Rect(rect.topleft, (point[0] - rect.left, point[1] - rect.top)),
		pygame.Rect(rect.midtop, (rect.right - point[0], point[1] - rect.top)),
		pygame.Rect(rect.midleft, (point[0] - rect.left, rect.bottom - point[1])),
		pygame.Rect(rect.center, (rect.right - point[0], rect.bottom - point[1]))
	]


class QuadTreeR:
	"""Implements a QuadTree containing arbitrary rects
	We number the quadrants like so:
	    ^
	  3 | 4
	x --+-->
	  1 | 2
	    y
	"""
	def __init__(self, bbox):
		self.bbox = pygame.Rect(bbox)
		self.rects = []
		self.subregions = None


	def add(self, rect, obj):
		if not self.subregions:
			self.rects.append((pygame.Rect(rect), obj))
			self.subregions = [QuadTreeR(r) for r in splitrect(self.bbox, self.bbox.center)]
		else:
			for s in self.subregions:
				if s.bbox.contains(rect):
					s.add(rect, obj)
					break
			else:
				self.rects.append((pygame.Rect(rect), obj))
	

	def remove(self, obj):
		for (r, o) in self.rects:
			if o == obj:
				self.rects.remove((r, o))
				return
		if self.subregions:
			for s in self.subregions:
				s.remove(obj)


	def __iter__(self):
		for (r, o) in self.rects:
			yield o
		if self.subregions:
			for s in self.subregions:
				for o in s:
					yield o


	def collideRectWalk(self, rect):
		"""generator that walks through the QuadTree, returning rects
		that collide with the given rect"""
		for (r, o) in self.rects:
			if r.colliderect(rect):
				yield o
		if self.subregions:
			for s in self.subregions:
				if s.bbox.colliderect(rect):
					for o in s.collideRectWalk(rect):
						yield o
	

	def collideRect(self, rect):
		return [o for o in self.collideRectWalk(rect)]


	def collidePointWalk(self, point):
		"""generator that walks through the QuadTree, returning rects
		that collide with the given point"""
		for (r, o) in self.rects:
			if r.collidepoint(point):
				yield o
		if self.subregions:
			for s in self.subregions:
				if s.bbox.collidepoint(point):
					for o in s.collidePointWalk(point):
						yield o


	def collidePoint(self, point):
		return [o for o in self.collidePointWalk(point)]


	def draw(self, rect = None):
		glColor3f(1.0, 1.0, 1.0)
		glBegin(GL_LINE_LOOP)
		glVertex3f(self.bbox.left, self.bbox.top, 0)
		glVertex3f(self.bbox.right, self.bbox.top, 0)
		glVertex3f(self.bbox.right, self.bbox.bottom, 0)
		glVertex3f(self.bbox.left, self.bbox.bottom, 0)
		glEnd()

		for (r, o) in self.rects:
			if rect and r.colliderect(rect):
				glColor3f(1.0, 1.0, 0.0)
			else:
				glColor3f(0.0, 0.0, 1.0)
			glBegin(GL_LINE_LOOP)
			glVertex3f(r.left, r.top, 0)
			glVertex3f(r.right, r.top, 0)
			glVertex3f(r.right, r.bottom, 0)
			glVertex3f(r.left, r.bottom, 0)
			glEnd()

		if self.subregions:
			for s in self.subregions:
				s.draw(rect)

		if rect:
			glColor3f(1.0, 0.0, 0.0)
			glBegin(GL_LINE_LOOP)
			glVertex3f(rect.left, rect.top, 0)
			glVertex3f(rect.right, rect.top, 0)
			glVertex3f(rect.right, rect.bottom, 0)
			glVertex3f(rect.left, rect.bottom, 0)
			glEnd()