#!/usr/bin/python
# encoding: utf-8

import numpy as np

__all__ = ["KDTree"]

debug = False

def square_distance(pointA, pointB):
    # squared euclidean distance
    A = pointA[1:]
    B = pointB
    return np.sum((A-B)**2)

class KDTreeNode():
    def __init__(self, point, left, right):
        self.point = point
        self.left = left
        self.right = right
    
    def is_leaf(self):
        return (self.left == None and self.right == None)

class KDTreeNeighbours():
    # Internal structure used in nearest-neighbours search.
    def __init__(self,query_point,eps):
        self.query_point = query_point
        self.eps = eps # squared
        self.current_best = []

    def add(self, point):
        sd = square_distance(point, self.query_point)
        if debug:
            print "sd: " + str(sd) + " | point: " + str(point)
 
        if sd <= self.eps: 
            if debug:
                print "\t ---- Yes"
            self.current_best.append([point])
        return
    
    def get_best(self):
        return [element[0] for element in self.current_best]
        
class KDTree():
    def __init__(self, data):
        def build_kdtree(point_list, depth):
            # code based on wikipedia article: http://en.wikipedia.org/wiki/Kd-tree
            if not point_list.any():
                return None

            # select axis based on depth so that axis cycles through all valid values
            k = (len(point_list[0]) -1)
            axis = (depth % k) + 1  # assumes all points have the same dimension

            # sort point list and choose median as pivot point,
            point_list = point_list[point_list[:,axis].argsort()]
            median = len(point_list)/2 # choose median
            if debug:
                print "Points: -------------" 
                print "Depth: " + str(depth)
                print "Length: " + str(len(point_list[0]))
                print "axis: " + str(axis)
                print "Median: " + str(median)
                print str(point_list)
                print "---------------------"

            # create node and recursively construct subtrees
            node = KDTreeNode(point=point_list[median],
                              left=build_kdtree(point_list[0:median], depth+1),
                              right=build_kdtree(point_list[median+1:], depth+1))
            return node
        
        self.root_node = build_kdtree(data, depth=0)
    
    @staticmethod
    def construct_from_data(data):
        tree = KDTree(data)
        return tree

    def query(self, query_point, eps):
        statistics = {'nodes_visited': 0, 'far_search': 0, 'leafs_reached': 0}
        
        #def nn_search(node, query_point, t, depth, best_neighbours):
        def nn_search(node, query_point,depth, best_neighbours):
            if node == None:
                return
            
            statistics['nodes_visited'] += 1
            
            # if we have reached a leaf, let's add to current best neighbours,
            # (if it's better than the worst one or if there is not enough neighbours)
            if node.is_leaf():
                statistics['leafs_reached'] += 1
                best_neighbours.add(np.array(node.point))
                return
            
            # select dimension for comparison (based on current depth)
            #axis = depth % len(query_point)
            k = len(query_point)
            axis = (depth % k) + 1 
            
            # figure out which subtree to search
            near_subtree = None # near subtree
            far_subtree = None # far subtree (perhaps we'll have to traverse it as well)
            
            # compare query_point and point of current node in selected dimension
            # and figure out which subtree is farther than the other
            if query_point[axis-1] < node.point[axis]:
                near_subtree = node.left
                far_subtree = node.right
            else:
                near_subtree = node.right
                far_subtree = node.left

            # recursively search through the tree until a leaf is found
            #nn_search(near_subtree, query_point, t, depth+1, best_neighbours)
            nn_search(near_subtree, np.array(query_point), depth+1, best_neighbours)

            # while unwinding the recursion, check if the current node
            # is closer to query point than the current best,
            # also, until t points have been found, search radius is infinity
            best_neighbours.add(np.array(node.point))
            
            # check whether there could be any points on the other side of the
            # splitting plane that are closer to the query point than the current best
            if (node.point[axis] - query_point[axis-1])**2 < best_neighbours.eps:
                if debug:
                    print "Distance: " + str((node.point[axis] - query_point[axis])**2)
                statistics['far_search'] += 1
                #nn_search(far_subtree, query_point, t, depth+1, best_neighbours)
                nn_search(far_subtree, np.array(query_point), depth+1, best_neighbours)            
 
            return
        
        # if there's no tree, there's no neighbors
        if self.root_node != None:
            #neighbours = KDTreeNeighbours(query_point, t)
            neighbours = KDTreeNeighbours(np.array(query_point), eps)
            #nn_search(self.root_node, query_point, t, depth=0, best_neighbours=neighbours)
            nn_search(self.root_node, np.array(query_point), depth=0, best_neighbours=neighbours)
            result = neighbours.get_best()
        else:
            result = []
       
        if debug: 
            print statistics
        return result
