#!/usr/bin/env python26
# -*- encoding: utf-8 -*-

"""
Quadtree intersection collision detection.

I do not recommend using this in a production environment _at all_.
"""

__author__ = "rfw"

import Image
from time import clock

DEPTH = 4

class QuadTree(object):
    def __init__(self, root, origin):
        self.root = root
        self.origin = origin

    def collides(self, other):
        return self.root.collides(self.origin, other.root, other.origin)

class QuadTreeNode(object):
    """
    An example quadtree node class.
    """
    def __init__(
        self,
        collidable,
        size,
        nw=None,
        ne=None,
        se=None,
        sw=None
    ):

        self.collidable = collidable

        self.size = size

        self.nw = nw
        self.ne = ne
        self.se = se
        self.sw = sw


    childless = property(lambda self: not (
        self.nw or
        self.ne or
        self.se or
        self.sw
    ))

    def __iter__(self):
        """
        Iterate through the quadtree clockwise, starting with northwewst.
        """
        if self.nw: yield self.nw
        if self.ne: yield self.ne
        if self.se: yield self.se
        if self.sw: yield self.sw

    def deep_collidable(self):
        """
        Recursively check if the node is collidable at any level.

        Automatically memoizes for efficiency.
        """
        if self.childless:
            return self.collidable

        if not hasattr(self, "_memoized_deep_collidable"):
            # calculate it
            for child in self:
                if child.deep_collidable():
                    self._memoized_deep_collidable = True
                    break
                else:
                    self._memoized_deep_collidable = False

        return self._memoized_deep_collidable

    def aabb_collides(self, origin, other, otherorigin):
        """
        Perform an AABB collision test to ensure that the quadrant is in the
        right place.
        """
        if origin[0] + self.size[0] < otherorigin[0]: return False
        if origin[0] > otherorigin[0] + other.size[0]: return False

        if origin[1] + self.size[1] < otherorigin[1]: return False
        if origin[1] > otherorigin[1] + other.size[1]: return False

        return True

    def collides(self, origin, other, otherorigin):
        # basic AABB pass
        if not self.aabb_collides(origin, other, otherorigin):
            return False

        # one or both of us have no kids
        if self.childless or other.childless:
            return self.deep_collidable() & other.deep_collidable()

        # dual quadtree traversal pass (XXX: messy!)
        for node in self:
            if node is self.nw:
                nodepos = (
                    origin[0],
                    origin[1]
                )
            elif node is self.ne:
                nodepos = (
                    origin[0] + node.size[0],
                    origin[1]
                )
            elif node is self.se:
                nodepos = (
                    origin[0] + node.size[0],
                    origin[1] + node.size[1]
                )
            elif node is self.sw:
                nodepos = (
                    origin[0],
                    origin[1] + node.size[1]
                )

            for othernode in other:
                if othernode is other.nw:
                    othernodepos = (
                        otherorigin[0],
                        otherorigin[1]
                    )
                elif node is other.ne:
                    othernodepos = (
                        otherorigin[0] + othernode.size[0],
                        otherorigin[1]
                    )
                elif othernode is other.se:
                    othernodepos = (
                        otherorigin[0] + othernode.size[0],
                        otherorigin[1] + othernode.size[1]
                    )
                elif node is other.sw:
                    othernodepos = (
                        otherorigin[0],
                        otherorigin[1] + othernode.size[1]
                    )

                if node.aabb_collides(nodepos, othernode, othernodepos) and \
                    node.deep_collidable() and \
                    othernode.deep_collidable():
                    return node.collides(nodepos, othernode, othernodepos)

        return False

def _naive_optimize_quadtree(node):
    """
    Naïvely optimize a quadtree via traversal.

    Returns the deepest nodes.
    """
    nodes = []

    if node.collidable: # no optimization required
        return [ node ]

    if node.nw and \
       node.ne and \
       node.se and \
       node.sw and \
       node.nw.collidable is True and \
       node.ne.collidable is True and \
       node.se.collidable is True and \
       node.sw.collidable is True:
        node.collidable = True

        node.nw = None
        node.ne = None
        node.se = None
        node.sw = None

        return [ node ]

    if node.nw:
        nodes.extend(_naive_optimize_quadtree(node.nw))

    if node.ne:
        nodes.extend(_naive_optimize_quadtree(node.ne))

    if node.se:
        nodes.extend(_naive_optimize_quadtree(node.se))

    if node.sw:
        nodes.extend(_naive_optimize_quadtree(node.sw))

    return nodes

def _naive_populate_quadtree(node, depth, origin):
    """
    Recursively populate a quadtree, naïvely (do not benchmark!)

    It will also attach position metadata which is meaningless in most cases.

    Returns the deepest nodes.
    """
    nodes = []

    node.position = origin

    if depth != 0:
        halfsize = (
            node.size[0] / 2.0,
            node.size[1] / 2.0
        )

        # northwest
        node.nw = QuadTreeNode(
            False,
            halfsize
        )
        nodes.extend(
            _naive_populate_quadtree(
                node.nw,
                depth - 1,
                (
                    origin[0],
                    origin[1]
                )
            )
        )

        # northeast
        node.ne = QuadTreeNode(
            False,
            halfsize
        )
        nodes.extend(
            _naive_populate_quadtree(
                node.ne,
                depth - 1,
                (
                    origin[0] + halfsize[0],
                    origin[1]
                )
            )
        )

        # southeast
        node.se = QuadTreeNode(
            False,
            halfsize
        )
        nodes.extend(
            _naive_populate_quadtree(
                node.se,
                depth - 1,
                (
                    origin[0] + halfsize[0],
                    origin[1] + halfsize[1]
                )
            )
        )

        # southwest
        node.sw = QuadTreeNode(
            False,
            halfsize
        )
        nodes.extend(
            _naive_populate_quadtree(
                node.sw,
                depth - 1,
                (
                    origin[0],
                    origin[1] + halfsize[1]
                )
            )
        )
    else:
        return [ node ]

    return nodes

def _construct_demo_quadtree(origin, image):
    """
    Construct a demo quadtree from an image (do not benchmark!)
    """
    im = Image.open(image)

    # construct a quadtree (very) naïvely

    starttime = clock()
    print "Populating quadtree...",
    root = QuadTreeNode(False, im.size)
    deepest = _naive_populate_quadtree(root, DEPTH, origin)
    quadtree = QuadTree(root, origin)
    print "done (%d deepest nodes, took %.2f seconds)." % (
        len(deepest),
        clock() - starttime
    )

    starttime = clock()
    print "Coloring quadtree...",
    collidable = set([])
    for i, color in enumerate(im.getdata()):
        x = i % im.size[0]
        y = i // im.size[0]

        if color != (0, 0, 0, 0):
            for node in deepest:
                relpos = (
                    node.position[0] - origin[0],
                    node.position[1] - origin   [1]
                )
                if relpos[0] < x < relpos[0] + node.size[0] and \
                   relpos[1] < y < relpos[1] + node.size[1]:
                    node.collidable = True
                    collidable.add(node)
                    break
    print "done (%d nodes set to collidable, took %.2f seconds)." % (
        len(collidable),
        clock() - starttime
    )

    starttime = clock()
    # optimize the quadtree
    print "Optimizing quadtree...",
    deepest = _naive_optimize_quadtree(quadtree.root)
    print "done (%d deepest nodes, took %.2f seconds)." % (
        len(deepest),
        clock() - starttime
    )

    print ""

    return quadtree

def main():
    tree = _construct_demo_quadtree((0, 0), "test.png")
    tree2 = QuadTree(tree.root, (-10, -10))
    tree3 = QuadTree(tree.root, (1000, 1000))

    starttime = clock()
    print "Performing 1000 true collisions...",
    for i in xrange(1000):
        assert tree.collides(tree2)
    print "done (took %.2f seconds)." % (clock() - starttime)

    starttime = clock()
    print "Performing 1000 false collisions...",
    for i in xrange(1000):
        assert not tree.collides(tree3)
    print "done (took %.2f seconds)." % (clock() - starttime)

if __name__ == "__main__":
    main()