#!python
#cython: boundscheck=False
#cython: wraparound=False
#cython: cdivision=True

# By Jake Vanderplas (2013) <jakevdp@cs.washington.edu>
# written for the scikit-learn project
# License: BSD

__all__ = ['KDTree']

DOC_DICT = {'BinaryTree': 'KDTree', 'binary_tree': 'kd_tree'}

VALID_METRICS = ['EuclideanDistance', 'ManhattanDistance',
                 'ChebyshevDistance', 'MinkowskiDistance']


include "binary_tree.pxi"

# Inherit KDTree from BinaryTree
cdef class KDTree(BinaryTree):
    __doc__ = CLASS_DOC.format(**DOC_DICT)
    pass


#----------------------------------------------------------------------
# The functions below specialized the Binary Tree as a KD Tree
#
#   Note that these functions use the concept of "reduced distance".
#   The reduced distance, defined for some metrics, is a quantity which
#   is more efficient to compute than the distance, but preserves the
#   relative rankings of the true distance.  For example, the reduced
#   distance for the Euclidean metric is the squared-euclidean distance.
#   For some metrics, the reduced distance is simply the distance.


cdef int allocate_data(BinaryTree tree, ITYPE_t n_nodes,
                       ITYPE_t n_features) except -1:
    """Allocate arrays needed for the KD Tree"""
    tree.node_bounds_arr = np.zeros((2, n_nodes, n_features), dtype=DTYPE)
    tree.node_bounds = get_memview_DTYPE_3D(tree.node_bounds_arr)
    return 0


cdef int init_node(BinaryTree tree, ITYPE_t i_node,
                   ITYPE_t idx_start, ITYPE_t idx_end) except -1:
    """Initialize the node for the dataset stored in tree.data"""
    cdef ITYPE_t n_features = tree.data.shape[1]
    cdef ITYPE_t i, j
    cdef DTYPE_t rad = 0

    cdef DTYPE_t* lower_bounds = &tree.node_bounds[0, i_node, 0]
    cdef DTYPE_t* upper_bounds = &tree.node_bounds[1, i_node, 0]
    cdef DTYPE_t* data = &tree.data[0, 0]
    cdef ITYPE_t* idx_array = &tree.idx_array[0]

    cdef DTYPE_t* data_row

    # determine Node bounds
    for j in range(n_features):
        lower_bounds[j] = INF
        upper_bounds[j] = -INF

    # Compute the actual data range.  At build time, this is slightly
    # slower than using the previously-computed bounds of the parent node,
    # but leads to more compact trees and thus faster queries.
    for i in range(idx_start, idx_end):
        data_row = data + idx_array[i] * n_features
        for j in range(n_features):
            lower_bounds[j] = fmin(lower_bounds[j], data_row[j])
            upper_bounds[j] = fmax(upper_bounds[j], data_row[j])

    for j in range(n_features):
        if tree.dist_metric.p == INF:
            rad = fmax(rad, 0.5 * (upper_bounds[j] - lower_bounds[j]))
        else:
            rad += pow(0.5 * abs(upper_bounds[j] - lower_bounds[j]),
                       tree.dist_metric.p)

    tree.node_data[i_node].idx_start = idx_start
    tree.node_data[i_node].idx_end = idx_end

    # The radius will hold the size of the circumscribed hypersphere measured
    # with the specified metric: in querying, this is used as a measure of the
    # size of each node when deciding which nodes to split.
    tree.node_data[i_node].radius = pow(rad, 1. / tree.dist_metric.p)
    return 0


cdef DTYPE_t min_rdist(BinaryTree tree, ITYPE_t i_node,
                       DTYPE_t* pt) nogil except -1:
    """Compute the minimum reduced-distance between a point and a node"""
    cdef ITYPE_t n_features = tree.data.shape[1]
    cdef DTYPE_t d, d_lo, d_hi, rdist=0.0
    cdef ITYPE_t j

    if tree.dist_metric.p == INF:
        for j in range(n_features):
            d_lo = tree.node_bounds[0, i_node, j] - pt[j]
            d_hi = pt[j] - tree.node_bounds[1, i_node, j]
            d = (d_lo + fabs(d_lo)) + (d_hi + fabs(d_hi))
            rdist = fmax(rdist, 0.5 * d)
    else:
        # here we'll use the fact that x + abs(x) = 2 * max(x, 0)
        for j in range(n_features):
            d_lo = tree.node_bounds[0, i_node, j] - pt[j]
            d_hi = pt[j] - tree.node_bounds[1, i_node, j]
            d = (d_lo + fabs(d_lo)) + (d_hi + fabs(d_hi))
            rdist += pow(0.5 * d, tree.dist_metric.p)

    return rdist


cdef DTYPE_t min_dist(BinaryTree tree, ITYPE_t i_node, DTYPE_t* pt) except -1:
    """Compute the minimum distance between a point and a node"""
    if tree.dist_metric.p == INF:
        return min_rdist(tree, i_node, pt)
    else:
        return pow(min_rdist(tree, i_node, pt), 1. / tree.dist_metric.p)


cdef DTYPE_t max_rdist(BinaryTree tree,
                       ITYPE_t i_node, DTYPE_t* pt) except -1:
    """Compute the maximum reduced-distance between a point and a node"""
    cdef ITYPE_t n_features = tree.data.shape[1]

    cdef DTYPE_t d, d_lo, d_hi, rdist=0.0
    cdef ITYPE_t j

    if tree.dist_metric.p == INF:
        for j in range(n_features):
            rdist = fmax(rdist, fabs(pt[j] - tree.node_bounds[0, i_node, j]))
            rdist = fmax(rdist, fabs(pt[j] - tree.node_bounds[1, i_node, j]))
    else:
        for j in range(n_features):
            d_lo = fabs(pt[j] - tree.node_bounds[0, i_node, j])
            d_hi = fabs(pt[j] - tree.node_bounds[1, i_node, j])
            rdist += pow(fmax(d_lo, d_hi), tree.dist_metric.p)

    return rdist


cdef DTYPE_t max_dist(BinaryTree tree, ITYPE_t i_node, DTYPE_t* pt) except -1:
    """Compute the maximum distance between a point and a node"""
    if tree.dist_metric.p == INF:
        return max_rdist(tree, i_node, pt)
    else:
        return pow(max_rdist(tree, i_node, pt), 1. / tree.dist_metric.p)


cdef inline int min_max_dist(BinaryTree tree, ITYPE_t i_node, DTYPE_t* pt,
                             DTYPE_t* min_dist, DTYPE_t* max_dist) nogil except -1:
    """Compute the minimum and maximum distance between a point and a node"""
    cdef ITYPE_t n_features = tree.data.shape[1]

    cdef DTYPE_t d, d_lo, d_hi
    cdef ITYPE_t j

    min_dist[0] = 0.0
    max_dist[0] = 0.0

    if tree.dist_metric.p == INF:
        for j in range(n_features):
            d_lo = tree.node_bounds[0, i_node, j] - pt[j]
            d_hi = pt[j] - tree.node_bounds[1, i_node, j]
            d = (d_lo + fabs(d_lo)) + (d_hi + fabs(d_hi))
            min_dist[0] = fmax(min_dist[0], 0.5 * d)
            max_dist[0] = fmax(max_dist[0],
                               fabs(pt[j] - tree.node_bounds[0, i_node, j]))
            max_dist[0] = fmax(max_dist[0],
                               fabs(pt[j] - tree.node_bounds[1, i_node, j]))
    else:
        # as above, use the fact that x + abs(x) = 2 * max(x, 0)
        for j in range(n_features):
            d_lo = tree.node_bounds[0, i_node, j] - pt[j]
            d_hi = pt[j] - tree.node_bounds[1, i_node, j]
            d = (d_lo + fabs(d_lo)) + (d_hi + fabs(d_hi))
            min_dist[0] += pow(0.5 * d, tree.dist_metric.p)
            max_dist[0] += pow(fmax(fabs(d_lo), fabs(d_hi)),
                               tree.dist_metric.p)

        min_dist[0] = pow(min_dist[0], 1. / tree.dist_metric.p)
        max_dist[0] = pow(max_dist[0], 1. / tree.dist_metric.p)

    return 0


cdef inline DTYPE_t min_rdist_dual(BinaryTree tree1, ITYPE_t i_node1,
                                   BinaryTree tree2, ITYPE_t i_node2) except -1:
    """Compute the minimum reduced distance between two nodes"""
    cdef ITYPE_t n_features = tree1.data.shape[1]

    cdef DTYPE_t d, d1, d2, rdist=0.0
    cdef DTYPE_t zero = 0.0
    cdef ITYPE_t j

    if tree1.dist_metric.p == INF:
        for j in range(n_features):
            d1 = (tree1.node_bounds[0, i_node1, j]
                  - tree2.node_bounds[1, i_node2, j])
            d2 = (tree2.node_bounds[0, i_node2, j]
                  - tree1.node_bounds[1, i_node1, j])
            d = (d1 + fabs(d1)) + (d2 + fabs(d2))

            rdist = fmax(rdist, 0.5 * d)
    else:
        # here we'll use the fact that x + abs(x) = 2 * max(x, 0)
        for j in range(n_features):
            d1 = (tree1.node_bounds[0, i_node1, j]
                  - tree2.node_bounds[1, i_node2, j])
            d2 = (tree2.node_bounds[0, i_node2, j]
                  - tree1.node_bounds[1, i_node1, j])
            d = (d1 + fabs(d1)) + (d2 + fabs(d2))

            rdist += pow(0.5 * d, tree1.dist_metric.p)

    return rdist


cdef inline DTYPE_t min_dist_dual(BinaryTree tree1, ITYPE_t i_node1,
                                  BinaryTree tree2, ITYPE_t i_node2) except -1:
    """Compute the minimum distance between two nodes"""
    return tree1.dist_metric._rdist_to_dist(min_rdist_dual(tree1, i_node1,
                                                           tree2, i_node2))


cdef inline DTYPE_t max_rdist_dual(BinaryTree tree1, ITYPE_t i_node1,
                                   BinaryTree tree2, ITYPE_t i_node2) except -1:
    """Compute the maximum reduced distance between two nodes"""
    cdef ITYPE_t n_features = tree1.data.shape[1]

    cdef DTYPE_t d, d1, d2, rdist=0.0
    cdef DTYPE_t zero = 0.0
    cdef ITYPE_t j

    if tree1.dist_metric.p == INF:
        for j in range(n_features):
            rdist = fmax(rdist, fabs(tree1.node_bounds[0, i_node1, j]
                                     - tree2.node_bounds[1, i_node2, j]))
            rdist = fmax(rdist, fabs(tree1.node_bounds[1, i_node1, j]
                                     - tree2.node_bounds[0, i_node2, j]))
    else:
        for j in range(n_features):
            d1 = fabs(tree1.node_bounds[0, i_node1, j]
                      - tree2.node_bounds[1, i_node2, j])
            d2 = fabs(tree1.node_bounds[1, i_node1, j]
                      - tree2.node_bounds[0, i_node2, j])
            rdist += pow(fmax(d1, d2), tree1.dist_metric.p)

    return rdist


cdef inline DTYPE_t max_dist_dual(BinaryTree tree1, ITYPE_t i_node1,
                                  BinaryTree tree2, ITYPE_t i_node2) except -1:
    """Compute the maximum distance between two nodes"""
    return tree1.dist_metric._rdist_to_dist(max_rdist_dual(tree1, i_node1,
                                                           tree2, i_node2))


######################################################################
# kd-forest !

import numpy as np
import scipy as sp

from cpython cimport set

class KDForest():

    def __init__(self, data, n_trees=1, leaf_size=40, **kwargs):
        self.n_trees = n_trees
        self.data = data
        self.dist_metric = DistanceMetric.get_metric("minkowski", **kwargs)
        self._dist_comps = 0

        ortho_group = sp.stats.special_ortho_group

        self.rotations = []
        self.trees = []
        for t in range(n_trees):
            rotation = ortho_group.rvs(data.shape[1])
            self.rotations.append(rotation)

            data_rot = np.dot(data, rotation)
            tree = KDTree(data_rot, tree_id=t, leaf_size=leaf_size)
            self.trees.append(tree)

    def query(self, X, k=1, sort_results=True, max_dist_comps=-1):
        dim = self.data.shape[1]
        X_rot = np.zeros((self.n_trees, dim))

        for t in range(self.n_trees):
            rotation = self.rotations[t]
            X_rot[t,:] = np.dot(X, rotation)

        cdef DTYPE_t[:, ::1] Xarr = get_memview_DTYPE_2D(X_rot)
        cdef DTYPE_t* pt = &Xarr[0, 0]
        cdef ITYPE_t pt_offset = 0

        cdef NeighborsHeap heap = NeighborsHeap(1, k)
        cdef NodeHeap nodeheap = NodeHeap(0)
        cdef NodeHeapData_t root
        cdef set visited = set()

        dist_comps = 0

        for t in range(self.n_trees):
            tree = self.trees[t]
            pt_offset = t * dim

            root.tree_id = t
            root.val = min_rdist(tree, 0, pt + pt_offset)
            root.i1 = 0
            root.i2 = 0

            nodeheap.push(root)

        dist_comps += self.n_trees

        while nodeheap.n > 0 and (max_dist_comps < 0 or dist_comps < max_dist_comps):
            nodeheap_item = nodeheap.pop()
            t = nodeheap_item.tree_id
            tree = self.trees[t]
            pt_offset = t * dim

            KDTree.visit_node(tree, pt + pt_offset, heap, nodeheap, nodeheap_item, visited)

            dist_comps += tree.get_n_calls() + tree.get_tree_stats()[2]
            tree.reset_n_calls()

        self._dist_comps = dist_comps

        distances, indices = heap.get_arrays(sort=sort_results)
        distances = self.dist_metric.rdist_to_dist(distances)
        return distances, indices

    def get_dist_comps(self):
        return self._dist_comps
