# Copyright Anne M. Archibald 2008
# Additional contributions by Patrick Varilly and Sturla Molden 2012
# Revision by Sturla Molden 2015
# Balanced kd-tree construction written by Jake Vanderplas for scikit-learn
# Released under the scipy license

# distutils: language = c++

from __future__ import absolute_import

import numpy as np
import scipy.sparse

cimport numpy as np
from numpy.math cimport INFINITY
    
from cpython.mem cimport PyMem_Malloc, PyMem_Realloc, PyMem_Free
from libc.string cimport memset, memcpy

cimport cython

from multiprocessing import cpu_count
import threading

cdef extern from "limits.h":
    long LONG_MAX

cdef extern from "ckdtree_methods.h":
    int number_of_processors
    
number_of_processors = cpu_count()

from libcpp.vector cimport vector
from libc cimport string

__all__ = ['cKDTree']

    
# Borrowed references
# ===================

cdef extern from *:
    
    struct ckdtree:
        pass
        
    int NPY_LIKELY(int)
    int NPY_UNLIKELY(int)

       
cdef extern from "ckdtree_decl.h":
    struct ckdtreenode:
        np.intp_t split_dim
        np.intp_t children
        np.float64_t split
        np.intp_t start_idx
        np.intp_t end_idx
        ckdtreenode *less
        ckdtreenode *greater
        np.intp_t _less
        np.intp_t _greater
    
    
# C++ helper functions
# ====================

cdef extern from "coo_entries.h":

    struct coo_entry:
        np.intp_t i
        np.intp_t j
        np.float64_t v

cdef extern from "ordered_pair.h":

    struct ordered_pair:
        np.intp_t i
        np.intp_t j

def new_object(obj):
    return obj.__new__(obj)
 
cdef extern from "cpp_utils.h": 
    object pickle_tree_buffer(vector[ckdtreenode] *buf)    
    object unpickle_tree_buffer(vector[ckdtreenode] *buf, object src)
    ckdtreenode *tree_buffer_root(vector[ckdtreenode] *buf)
    ordered_pair *ordered_pair_vector_buf(vector[ordered_pair] *buf)
    coo_entry *coo_entry_vector_buf(vector[coo_entry] *buf)
    void *tree_buffer_pointer(vector[ckdtreenode] *buf)
    np.intp_t *npy_intp_vector_buf(vector[np.intp_t] *buf)
    np.float64_t *npy_float64_vector_buf(vector[np.float64_t] *buf)
    ctypedef void *intvector_ptr_t 



# coo_entry wrapper
# =================

cdef class coo_entries:

    cdef: 
        readonly object __array_interface__
        vector[coo_entry] *buf
        
    def __cinit__(coo_entries self):    
        self.buf = NULL

    def __init__(coo_entries self):    
        self.buf = new vector[coo_entry]()
        
    def __dealloc__(coo_entries self):
        if self.buf != NULL:
            del self.buf
            
    # The methods ndarray, dict, coo_matrix, and dok_matrix must only
    # be called after the buffer is filled with coo_entry data. This
    # is because std::vector can reallocate its internal buffer when
    # push_back is called.
            
    def ndarray(coo_entries self):    
        cdef: 
            coo_entry *pr
            np.uintp_t uintptr
            np.intp_t n
        _dtype = [('i',np.intp),('j',np.intp),('v',np.float64)]
        res_dtype = np.dtype(_dtype, align = True)
        n = <np.intp_t> self.buf.size()
        if NPY_LIKELY(n > 0):
            pr = coo_entry_vector_buf(self.buf) 
            uintptr = <np.uintp_t> (<void*> pr)
            dtype = np.dtype(np.uint8)               
            self.__array_interface__ = dict(
                data = (uintptr, False),
                descr = dtype.descr,
                shape = (n*sizeof(coo_entry),),
                strides = (dtype.itemsize,),
                typestr = dtype.str,
                version = 3,
            ) 
            return np.asarray(self).view(dtype=res_dtype)
        else:
            return np.empty(shape=(0,), dtype=res_dtype)
        
    def dict(coo_entries self):
        cdef:
            np.intp_t i, j, k, n
            np.float64_t v
            coo_entry *pr
            dict res_dict
        n = <np.intp_t> self.buf.size()
        if NPY_LIKELY(n > 0):
            pr = coo_entry_vector_buf(self.buf)        
            res_dict = dict()           
            for k in range(n):                    
                i = pr[k].i
                j = pr[k].j
                v = pr[k].v                    
                res_dict[(i,j)] = v
            return res_dict
        else:
            return {}
    
    def coo_matrix(coo_entries self, m, n):
        res_arr = self.ndarray()
        return scipy.sparse.coo_matrix(
                       (res_arr['v'], (res_arr['i'], res_arr['j'])),
                                       shape=(m, n))
        
    def dok_matrix(coo_entries self, m, n):
        return self.coo_matrix(m,n).todok()
        
        
# ordered_pair wrapper
# ====================

cdef class ordered_pairs:

    cdef: 
        readonly object __array_interface__
        vector[ordered_pair] *buf
        
    def __cinit__(ordered_pairs self):
        self.buf = NULL

    def __init__(ordered_pairs self):
        self.buf = new vector[ordered_pair]()
        
    def __dealloc__(ordered_pairs self):
        if self.buf != NULL:
            del self.buf
            
    # The methods ndarray and set must only be called after the buffer 
    # is filled with ordered_pair data.
    
    def ndarray(ordered_pairs self):
        cdef: 
            ordered_pair *pr
            np.uintp_t uintptr 
            np.intp_t n            
        n = <np.intp_t> self.buf.size()
        if NPY_LIKELY(n > 0):
            pr = ordered_pair_vector_buf(self.buf) 
            uintptr = <np.uintp_t> (<void*> pr)       
            dtype = np.dtype(np.intp)               
            self.__array_interface__ = dict(
                data = (uintptr, False),
                descr = dtype.descr,
                shape = (n,2),
                strides = (2*dtype.itemsize,dtype.itemsize),
                typestr = dtype.str,
                version = 3,
            )
            return np.asarray(self)
        else:
            return np.empty(shape=(0,2), dtype=np.intp)
        
    def set(ordered_pairs self):        
        cdef: 
            ordered_pair *pair
            np.intp_t i, n
            set results
        results = set()
        pair = ordered_pair_vector_buf(self.buf)
        n = <np.intp_t> self.buf.size()
        if sizeof(long) < sizeof(np.intp_t):
            # Needed for Python 2.x on Win64
            for i in range(n):
                results.add((int(pair.i), int(pair.j)))
                pair += 1 
        else:
            # other platforms
            for i in range(n):
                results.add((pair.i, pair.j))
                pair += 1
        return results
    
        
            
# Tree structure exposed to Python
# ================================

cdef class cKDTreeNode:
    """
    class cKDTreeNode

    This class exposes a Python view of a node in the cKDTree object.
    
    All attributes are read-only.
    
    Attributes
    ----------
    level : int
        The depth of the node. 0 is the level of the root node.
    split_dim : int
        The dimension along which this node is split. If this value is -1  
        the node is a leafnode in the kd-tree. Leafnodes are not split further
        and scanned by brute force.
    split : float
        The value used to separate split this node. Points with value >= split
        in the split_dim dimension are sorted to the 'greater' subnode 
        whereas those with value < split are sorted to the 'lesser' subnode.
    children : int
        The number of data points sorted to this node.
    data_points : ndarray of float64
        An array with the data points sorted to this node.
    indices : ndarray of intp
        An array with the indices of the data points sorted to this node. The
        indices refer to the position in the data set used to construct the
        kd-tree.
    lesser : cKDTreeNode or None
        Subnode with the 'lesser' data points. This attribute is None for
        leafnodes.
    greater : cKDTreeNode or None
        Subnode with the 'greater' data points. This attribute is None for
        leafnodes.

    """
    cdef:
        readonly np.intp_t    level
        readonly np.intp_t    split_dim
        readonly np.intp_t    children
        readonly np.float64_t split
        ckdtreenode           *_node
        np.ndarray            _data
        np.ndarray            _indices
        
    cdef void _setup(cKDTreeNode self):
        self.split_dim = self._node.split_dim
        self.children = self._node.children
        self.split = self._node.split
        
    property data_points:   
        def __get__(cKDTreeNode self):
            return self._data[self.indices,:]
                     
    property indices:
        def __get__(cKDTreeNode self):
            cdef np.intp_t i, start, stop
            if self.split_dim == -1:
                start = self._node.start_idx
                stop = self._node.end_idx
                return self._indices[start:stop]
            else:
                return np.hstack([self.lesser.indices, 
                           self.greater.indices])
                           
    property lesser:
        def __get__(cKDTreeNode self):
            if self.split_dim == -1:
                return None
            else:
                n = cKDTreeNode()
                n._node = self._node.less
                n._data = self._data
                n._indices = self._indices
                n.level = self.level + 1
                n._setup()
                return n
                
    property greater:
        def __get__(cKDTreeNode self):
            if self.split_dim == -1:
                return None
            else:
                n = cKDTreeNode()
                n._node = self._node.greater
                n._data = self._data
                n._indices = self._indices
                n.level = self.level + 1
                n._setup()
                return n

    
# Main cKDTree class
# ==================

cdef extern from "ckdtree_methods.h":

    # External build and query methods in C++. These will internally
    # release the GIL to avoid locking up the interpreter.
    
    int ckdtree_isinf(np.float64_t x)
    
    object build_ckdtree(ckdtree *self, 
                         np.intp_t start_idx, 
                         np.intp_t end_idx,
                         np.float64_t *maxes, 
                         np.float64_t *mins, 
                         int _median, 
                         int _compact)

    object build_weights(ckdtree *self, 
                         np.float64_t *node_weights,
                         np.float64_t *weights)
       
    object query_knn(const ckdtree *self, 
                     np.float64_t *dd, 
                     np.intp_t    *ii, 
                     const np.float64_t *xx,
                     const np.intp_t    n,
                     const np.intp_t    *k, 
                     const np.intp_t    nk, 
                     const np.intp_t    kmax, 
                     const np.float64_t eps, 
                     const np.float64_t p, 
                     const np.float64_t distance_upper_bound) 
                     
    object query_pairs(const ckdtree *self, 
                       const np.float64_t r, 
                       const np.float64_t p, 
                       const np.float64_t eps,
                       vector[ordered_pair] *results)

    object count_neighbors_unweighted(const ckdtree *self,
                           const ckdtree *other,
                           np.intp_t     n_queries,
                           np.float64_t  *real_r,
                           np.intp_t     *results,
                           const np.float64_t p,
                           int cumulative)

    object count_neighbors_weighted(const ckdtree *self,
                           const ckdtree *other,
                           np.float64_t  *self_weights,
                           np.float64_t  *other_weights,
                           np.float64_t  *self_node_weights,
                           np.float64_t  *other_node_weights,
                           np.intp_t     n_queries,
                           np.float64_t  *real_r,
                           np.float64_t     *results,
                           const np.float64_t p,
                           int cumulative)

    object query_ball_point(const ckdtree *self,
                            const np.float64_t *x,
                            const np.float64_t r,
                            const np.float64_t p,
                            const np.float64_t eps,
                            const np.intp_t n_queries,
                            vector[np.intp_t] **results)

    object query_ball_tree(const ckdtree *self,
                           const ckdtree *other,
                           const np.float64_t r,
                           const np.float64_t p,
                           const np.float64_t eps,
                           vector[np.intp_t] **results)                     
     
    object sparse_distance_matrix(const ckdtree *self,
                                  const ckdtree *other,
                                  const np.float64_t p,
                                  const np.float64_t max_distance,
                                  vector[coo_entry] *results)                    
                      
                      
cdef public class cKDTree [object ckdtree, type ckdtree_type]:
    """
    cKDTree(data, leafsize=16, compact_nodes=True, copy_data=False,
            balanced_tree=True)

    kd-tree for quick nearest-neighbor lookup

    This class provides an index into a set of k-dimensional points
    which can be used to rapidly look up the nearest neighbors of any
    point. 

    The algorithm used is described in Maneewongvatana and Mount 1999. 
    The general idea is that the kd-tree is a binary trie, each of whose
    nodes represents an axis-aligned hyperrectangle. Each node specifies
    an axis and splits the set of points based on whether their coordinate
    along that axis is greater than or less than a particular value. 

    During construction, the axis and splitting point are chosen by the 
    "sliding midpoint" rule, which ensures that the cells do not all
    become long and thin. 

    The tree can be queried for the r closest neighbors of any given point 
    (optionally returning only those within some maximum distance of the 
    point). It can also be queried, with a substantial gain in efficiency, 
    for the r approximate closest neighbors.

    For large dimensions (20 is already large) do not expect this to run 
    significantly faster than brute force. High-dimensional nearest-neighbor
    queries are a substantial open problem in computer science.

    Parameters
    ----------
    data : array_like, shape (n,m)
        The n data points of dimension m to be indexed. This array is 
        not copied unless this is necessary to produce a contiguous 
        array of doubles, and so modifying this data will result in 
        bogus results. The data are also copied if the kd-tree is built
        with copy_data=True.
    leafsize : positive int, optional
        The number of points at which the algorithm switches over to
        brute-force. Default: 16.
    compact_nodes : bool, optional    
        If True, the kd-tree is built to shrink the hyperrectangles to
        the actual data range. This usually gives a more compact tree that 
        is robust against degenerated input data and gives faster queries 
        at the expense of longer build time. Default: True.
    copy_data : bool, optional
        If True the data is always copied to protect the kd-tree against 
        data corruption. Default: False.
    balanced_tree : bool, optional    
        If True, the median is used to split the hyperrectangles instead of 
        the midpoint. This usually gives a more compact tree and 
        faster queries at the expense of longer build time. Default: True.
    boxsize : array_like or scalar, optional
        Apply a m-d toroidal topology to the KDTree.. The topology is generated 
        by :math:`x_i + n_i L_i` where :math:`n_i` are integers and :math:`L_i`
        is the boxsize along i-th dimension. The input data shall be wrapped 
        into :math:`[0, L_i)`. A ValueError is raised if any of the data is
        outside of this bound.

    Attributes
    ----------
    data : ndarray, shape (n,m)
        The n data points of dimension m to be indexed. This array is
        not copied unless this is necessary to produce a contiguous
        array of doubles. The data are also copied if the kd-tree is built
        with `copy_data=True`.
    leafsize : positive int
        The number of points at which the algorithm switches over to
        brute-force.
    m : int
        The dimension of a single data-point.
    n : int
        The number of data points.
    maxes : ndarray, shape (m,)
        The maximum value in each dimension of the n data points.
    mins : ndarray, shape (m,)
        The minimum value in each dimension of the n data points.
    tree : object, class cKDTreeNode
        This class exposes a Python view of the root node in the cKDTree object.
    size : int
        The number of nodes in the tree.

    See Also
    --------
    KDTree : Implementation of `cKDTree` in pure Python

    """
    cdef:
        vector[ckdtreenode]      *tree_buffer
        ckdtreenode              *ctree 
        readonly cKDTreeNode     tree 
        readonly np.ndarray      data
        np.float64_t             *raw_data
        readonly np.intp_t       n, m
        readonly np.intp_t       leafsize
        readonly np.ndarray      maxes
        np.float64_t             *raw_maxes
        readonly np.ndarray      mins
        np.float64_t             *raw_mins
        readonly np.ndarray      indices
        np.intp_t                *raw_indices
        np.ndarray               _median_workspace
        readonly object          boxsize
        np.ndarray               boxsize_data
        np.float64_t             *raw_boxsize_data
        readonly np.intp_t       size

    def __cinit__(cKDTree self):
        self.tree_buffer = NULL        
            
    def __init__(cKDTree self, data, np.intp_t leafsize=16, compact_nodes=True, 
            copy_data=False, balanced_tree=True, boxsize=None):
        cdef np.ndarray[np.float64_t, ndim=2] data_arr
        cdef np.float64_t *tmp
        cdef int _median, _compact
        cdef np.ndarray[np.float64_t, ndim=1] boxsize_arr
        data_arr = np.ascontiguousarray(data, dtype=np.float64)
        if copy_data and (data_arr is data):
            data_arr = data_arr.copy()
        self.data = data_arr
        self.n = data_arr.shape[0]
        self.m = data_arr.shape[1]
        self.leafsize = leafsize
        if self.leafsize<1:
            raise ValueError("leafsize must be at least 1")

        if boxsize is None:
            self.boxsize = None
            self.boxsize_data = None
        else:
            boxsize_arr = np.empty(2 * self.m, dtype=np.float64)
            boxsize_arr[:self.m] = boxsize
            boxsize_arr[self.m:] = 0.5 * boxsize_arr[:self.m]
            # FIXME: how to use a matching del if new is used?
            self.boxsize_data = boxsize_arr
            self.boxsize = boxsize_arr[:self.m].copy()
            periodic_mask = self.boxsize > 0
            if ((self.data >= self.boxsize[None, :])[:, periodic_mask]).any():
                raise ValueError("Some input data are greater than the size of the periodic box.")
            if ((self.data < 0)[:, periodic_mask]).any():
                raise ValueError("Negative input data are outside of the periodic box.")

        self.maxes = np.ascontiguousarray(np.amax(self.data,axis=0), dtype=np.float64)
        self.mins = np.ascontiguousarray(np.amin(self.data,axis=0), dtype=np.float64)
        self.indices = np.ascontiguousarray(np.arange(self.n,dtype=np.intp))

        self._pre_init()

        _compact = 1 if compact_nodes else 0
        _median = 1 if balanced_tree else 0
        if _median:
            self._median_workspace = np.zeros(self.n)

        self.tree_buffer = new vector[ckdtreenode]()
        
        try:
            tmp = <np.float64_t*> PyMem_Malloc(self.m*2*sizeof(np.float64_t))
            if tmp == NULL: raise MemoryError()            
            memcpy(tmp, self.raw_maxes, self.m*sizeof(np.float64_t))
            memcpy(tmp + self.m, self.raw_mins, self.m*sizeof(np.float64_t))
            build_ckdtree(<ckdtree*> self, 0, self.n, tmp, tmp + self.m, 
                _median, _compact)
        finally:
            PyMem_Free(tmp)

        self._median_workspace = None
        
        # set up the tree structure pointers
        self._post_init()
        
        # make the tree viewable from Python
        self.tree = cKDTreeNode()
        self.tree._node = self.ctree
        self.tree._data = self.data
        self.tree._indices = self.indices
        self.tree.level = 0
        self.tree._setup()

    cdef int _pre_init(cKDTree self) except -1:

        # finalize the pointers from array attributes

        self.raw_data = <np.float64_t*> np.PyArray_DATA(self.data)
        self.raw_maxes = <np.float64_t*> np.PyArray_DATA(self.maxes)
        self.raw_mins = <np.float64_t*> np.PyArray_DATA(self.mins)
        self.raw_indices = <np.intp_t*> np.PyArray_DATA(self.indices)

        if self.boxsize_data is not None:
            self.raw_boxsize_data = <np.float64_t*>np.PyArray_DATA(self.boxsize_data)

        return 0        

    cdef int _post_init(cKDTree self) except -1:
        # finalize the tree points, this calls _post_init_traverse
        
        self.ctree = tree_buffer_root(self.tree_buffer)

        # set the size attribute after tree_buffer is built
        self.size = self.tree_buffer.size()

        return self._post_init_traverse(self.ctree)
         
    cdef int _post_init_traverse(cKDTree self, ckdtreenode *node) except -1:
        # recurse the tree and re-initialize
        # "less" and "greater" fields
        if node.split_dim == -1:
            # leafnode
            node.less = NULL
            node.greater = NULL
        else:
            node.less = self.ctree + node._less
            node.greater = self.ctree + node._greater
            self._post_init_traverse(node.less)
            self._post_init_traverse(node.greater)
                
        return 0
        

    def __dealloc__(cKDTree self):
        if self.tree_buffer != NULL:
            del self.tree_buffer

    # -----
    # query
    # -----
    
    @cython.boundscheck(False)
    def query(cKDTree self, object x, object k=1, np.float64_t eps=0,
              np.float64_t p=2, np.float64_t distance_upper_bound=INFINITY,
              np.intp_t n_jobs=1):
        """
        query(self, x, k=1, eps=0, p=2, distance_upper_bound=np.inf, n_jobs=1)

        Query the kd-tree for nearest neighbors

        Parameters
        ----------
        x : array_like, last dimension self.m
            An array of points to query.
        k : list of integer or integer
            The list of k-th nearest neighbors to return. If k is an 
            integer it is treated as a list of [1, ... k] (range(1, k+1)).
            Note that the counting starts from 1.
        eps : non-negative float
            Return approximate nearest neighbors; the k-th returned value 
            is guaranteed to be no further than (1+eps) times the 
            distance to the real k-th nearest neighbor.
        p : float, 1<=p<=infinity
            Which Minkowski p-norm to use. 
            1 is the sum-of-absolute-values "Manhattan" distance
            2 is the usual Euclidean distance
            infinity is the maximum-coordinate-difference distance
        distance_upper_bound : nonnegative float
            Return only neighbors within this distance.  This is used to prune
            tree searches, so if you are doing a series of nearest-neighbor
            queries, it may help to supply the distance to the nearest neighbor
            of the most recent point.
        n_jobs : int, optional
            Number of jobs to schedule for parallel processing. If -1 is given
            all processors are used. Default: 1.
                        
        Returns
        -------
        d : array of floats
            The distances to the nearest neighbors. 
            If ``x`` has shape ``tuple+(self.m,)``, then ``d`` has shape ``tuple+(k,)``.
            When k == 1, the last dimension of the output is squeezed.
            Missing neighbors are indicated with infinite distances.
        i : ndarray of ints
            The locations of the neighbors in ``self.data``.
            If ``x`` has shape ``tuple+(self.m,)``, then ``i`` has shape ``tuple+(k,)``.
            When k == 1, the last dimension of the output is squeezed.
            Missing neighbors are indicated with ``self.n``.

        Notes
        -----
        If the KD-Tree is periodic, the position ``x`` is wrapped into the
        box.
        
        When the input k is a list, a query for arange(max(k)) is performed, but
        only columns that store the requested values of k are preserved. This is 
        implemented in a manner that reduces memory usage.

        Examples
        --------

        >>> import numpy as np
        >>> from scipy.spatial import cKDTree
        >>> x, y = np.mgrid[0:5, 2:8]
        >>> tree = cKDTree(np.c_[x.ravel(), y.ravel()])

        To query the nearest neighbours and return squeezed result, use

        >>> dd, ii = tree.query([[0, 0], [2.1, 2.9]], k=1)
        >>> print(dd, ii)
        [2.         0.14142136] [ 0 13]

        To query the nearest neighbours and return unsqueezed result, use

        >>> dd, ii = tree.query([[0, 0], [2.1, 2.9]], k=[1])
        >>> print(dd, ii)
        [[2.        ]
         [0.14142136]] [[ 0]
         [13]]

        To query the second nearest neighbours and return unsqueezed result, use

        >>> dd, ii = tree.query([[0, 0], [2.1, 2.9]], k=[2])
        >>> print(dd, ii)
        [[2.23606798]
         [0.90553851]] [[ 6]
         [12]]

        To query the first and second nearest neighbours, use

        >>> dd, ii = tree.query([[0, 0], [2.1, 2.9]], k=2)
        >>> print(dd, ii)
        [[2.         2.23606798]
         [0.14142136 0.90553851]] [[ 0  6]
         [13 12]]

        or, be more specific

        >>> dd, ii = tree.query([[0, 0], [2.1, 2.9]], k=[1, 2])
        >>> print(dd, ii)
        [[2.         2.23606798]
         [0.14142136 0.90553851]] [[ 0  6]
         [13 12]]

        """
        
        cdef:
            np.intp_t n, i, j
            int overflown
        
        x_arr = np.asarray(x, dtype=np.float64)
        if x_arr.ndim == 0 or x_arr.shape[x_arr.ndim - 1] != self.m:
            raise ValueError("x must consist of vectors of length %d but "
                             "has shape %s" % (int(self.m), np.shape(x)))
        if p < 1:
            raise ValueError("Only p-norms with 1<=p<=infinity permitted")
        if x_arr.ndim == 1:
            single = True
            x_arr = x_arr[np.newaxis,:]
        else:
            single = False

        nearest = False
        if np.isscalar(k):
            if k == 1:
                nearest = True
            k = np.arange(1, k + 1)
    
        retshape = np.shape(x)[:-1]
        n = <np.intp_t> np.prod(retshape)
        xx = np.ascontiguousarray(x_arr).reshape(n, self.m)

        # The C++ function touches all dd and ii entries,
        # setting the missing values.

        dd = np.empty((n,len(k)),dtype=np.float64)
        ii = np.empty((n,len(k)),dtype=np.intp)

        # Do the query in an external C++ function. 
        # The GIL will be released in the external query function.
        def _thread_func(self, np.intp_t start, np.intp_t stop):
            cdef: 
                np.ndarray[np.intp_t,ndim=2] _ii = ii
                np.ndarray[np.float64_t,ndim=2] _dd = dd
                np.ndarray[np.float64_t,ndim=2] _xx = xx
                np.ndarray[np.intp_t,ndim=1] _k = np.array(k, dtype=np.intp)
            
            kmax = np.max(k)

            query_knn(<ckdtree*>self, &_dd[start,0], &_ii[start,0], 
                &_xx[start,0], stop-start, &_k[0], len(k), kmax, eps, p, distance_upper_bound)
        
        if (n_jobs == -1): 
            n_jobs = number_of_processors
        
        if n_jobs > 1:
            # static scheduling without load balancing is good enough
                             
            ranges = [(j * n // n_jobs, (j + 1) * n // n_jobs)
                            for j in range(n_jobs)]
            
            # There might be n_jobs+1 threads spawned here, but only n_jobs of 
            # them will do significant work.
            threads = [threading.Thread(target=_thread_func,
                                args=(self, start, stop)) for start, stop in ranges]

            # Set the daemon flag so the process can be aborted, 
            # start all threads and wait for completion.
            for t in threads:
                t.daemon = True
                t.start()
            for t in threads: 
                t.join()
        else:
            _thread_func(self, 0, n)
                
        # massage the output in conformabity to the documented behavior

        if sizeof(long) < sizeof(np.intp_t):
            # ... e.g. Windows 64
            overflown = False
            for i in range(n):
                for j in range(len(k)):
                    if ii[i,j] > <np.intp_t>LONG_MAX:
                        # C long overlow, return array of dtype=np.int_p
                        overflown = True
                        break
                if overflown: 
                    break

            if overflown:
                ddret = np.reshape(dd,retshape+(len(k),))
                iiret = np.reshape(ii,retshape+(len(k),))
            else:
                ddret = np.reshape(dd,retshape+(len(k),))
                iiret = np.reshape(ii,retshape+(len(k),)).astype(int) 
                        
        else:
            # ... most other platforms
            ddret = np.reshape(dd,retshape+(len(k),))
            iiret = np.reshape(ii,retshape+(len(k),))

        if nearest:
            ddret = ddret[..., 0]
            iiret = iiret[..., 0]
            # the only case where we return a python scalar
            if single:
                ddret = float(ddret)
                iiret = int(iiret)
            
        return ddret, iiret

    # ----------------
    # query_ball_point
    # ----------------

    def query_ball_point(cKDTree self, object x, np.float64_t r,
                         np.float64_t p=2., np.float64_t eps=0, n_jobs=1):
        """
        query_ball_point(self, x, r, p=2., eps=0)
        
        Find all points within distance r of point(s) x.

        Parameters
        ----------
        x : array_like, shape tuple + (self.m,)
            The point or points to search for neighbors of.
        r : positive float
            The radius of points to return.
        p : float, optional
            Which Minkowski p-norm to use.  Should be in the range [1, inf].
        eps : nonnegative float, optional
            Approximate search. Branches of the tree are not explored if their
            nearest points are further than ``r / (1 + eps)``, and branches are
            added in bulk if their furthest points are nearer than
            ``r * (1 + eps)``.
        n_jobs : int, optional
            Number of jobs to schedule for parallel processing. If -1 is given
            all processors are used. Default: 1.

        Returns
        -------
        results : list or array of lists
            If `x` is a single point, returns a list of the indices of the
            neighbors of `x`. If `x` is an array of points, returns an object
            array of shape tuple containing lists of neighbors.

        Notes
        -----
        If you have many points whose neighbors you want to find, you may save
        substantial amounts of time by putting them in a cKDTree and using
        query_ball_tree.

        Examples
        --------
        >>> from scipy import spatial
        >>> x, y = np.mgrid[0:4, 0:4]
        >>> points = np.c_[x.ravel(), y.ravel()]
        >>> tree = spatial.cKDTree(points)
        >>> tree.query_ball_point([2, 0], 1)
        [4, 8, 9, 12]

        """
        
        cdef:
            np.ndarray[np.float64_t, ndim=1, mode="c"] xx
            np.ndarray[np.float64_t, ndim=2, mode="c"] vxx
            vector[np.intp_t] *vres
            vector[np.intp_t] **vvres
            np.uintp_t vvres_uintp
            np.intp_t *cur
            list tmp
            np.intp_t i, j, n, m
        
        vres = NULL
        vvres = NULL
        
        try:
               
            x = np.asarray(x, dtype=np.float64)
            if x.shape[-1] != self.m:
                raise ValueError("Searching for a %d-dimensional point in a "
                                 "%d-dimensional KDTree" % 
                                     (int(x.shape[-1]), int(self.m)))
            if len(x.shape) == 1:
                vres = new vector[np.intp_t]()
                xx = np.ascontiguousarray(x, dtype=np.float64)
                query_ball_point(<ckdtree*> self, &xx[0], r, p, eps, 1, &vres)
                n = <np.intp_t> vres.size()
                tmp = n * [None]
                if NPY_LIKELY(n > 0):
                    cur = npy_intp_vector_buf(vres)
                    for i in range(n):
                        tmp[i] = cur[0]
                        cur += 1
                result = tmp
            
            else:
                retshape = x.shape[:-1]
                
                # allocate an array of std::vector<npy_intp>
                n = np.prod(retshape)
                vvres = (<vector[np.intp_t] **> 
                    PyMem_Malloc(n * sizeof(intvector_ptr_t)))
                if vvres == NULL:
                    raise MemoryError()
                
                memset(<void*> vvres, 0, n * sizeof(intvector_ptr_t))      
            
                for i in range(n):
                    vvres[i] = new vector[np.intp_t]()
                
                result = np.empty(retshape, dtype=object)
                
                vxx = np.zeros((n,self.m), dtype=np.float64)
                i = 0
                for c in np.ndindex(retshape):
                    vxx[i,:] = x[c]
                    i += 1
                    
                # multithreading logic is similar to cKDTree.query
                                        
                if (n_jobs == -1): 
                    n_jobs = number_of_processors
        
                if n_jobs > 1:
                
                    CHUNK = n//n_jobs if n//n_jobs else n
                             
                    def _thread_func(self, _j, _vxx, r, p, eps, _vvres, CHUNK): 
                        cdef: 
                            np.intp_t j = _j
                            np.ndarray[np.float64_t,ndim=2] vxx = _vxx
                            vector[np.intp_t] **vvres                   
                            np.intp_t start = j*CHUNK
                            np.intp_t stop = start + CHUNK
                        stop = n if stop > n else stop
                        vvres = (<vector[np.intp_t] **> 
                                  (<void*> (<np.uintp_t> _vvres)))                                    
                        if start < n:
                            query_ball_point(<ckdtree*>self, &vxx[start,0], 
                                r, p, eps, stop-start, vvres+start)
                                
                    vvres_uintp = <np.uintp_t> (<void*> vvres)
                    threads = [threading.Thread(target=_thread_func,
                               args=(self, j, vxx, r, p, eps,vvres_uintp,CHUNK))
                                  for j in range(1+(n//CHUNK))]
                    for t in threads:
                        t.daemon = True
                        t.start()
                    for t in threads: 
                        t.join()
                                                                
                else:
                
                    query_ball_point(<ckdtree*>self, &vxx[0,0], r, p, eps, 
                        n, vvres)
                
                i = 0
                for c in np.ndindex(retshape):
                    m = <np.intp_t> (vvres[i].size())
                    if NPY_LIKELY(m > 0):
                        tmp = m * [None]
                        cur = npy_intp_vector_buf(vvres[i])
                        for j in range(m):
                            tmp[j] = cur[0]
                            cur += 1
                        result[c] = sorted(tmp)
                    else:
                        result[c] = []
                    i += 1
        
        finally:
            if vres != NULL: 
                del vres
                
            if vvres != NULL:
                for i in range(n):
                    if vvres[i] != NULL:
                        del vvres[i]     
                PyMem_Free(vvres)
                
        return result   
            

    # ---------------
    # query_ball_tree
    # ---------------
    
    def query_ball_tree(cKDTree self, cKDTree other,
                        np.float64_t r, np.float64_t p=2., np.float64_t eps=0):
        """
        query_ball_tree(self, other, r, p=2., eps=0)

        Find all pairs of points whose distance is at most r

        Parameters
        ----------
        other : cKDTree instance
            The tree containing points to search against.
        r : float
            The maximum distance, has to be positive.
        p : float, optional
            Which Minkowski norm to use.  `p` has to meet the condition
            ``1 <= p <= infinity``.
        eps : float, optional
            Approximate search.  Branches of the tree are not explored
            if their nearest points are further than ``r/(1+eps)``, and
            branches are added in bulk if their furthest points are nearer
            than ``r * (1+eps)``.  `eps` has to be non-negative.

        Returns
        -------
        results : list of lists
            For each element ``self.data[i]`` of this tree, ``results[i]`` is a
            list of the indices of its neighbors in ``other.data``.

        """
        
        cdef: 
            vector[np.intp_t] **vvres
            np.intp_t i, j, n, m
            np.intp_t *cur
            list results
            list tmp

        # Make sure trees are compatible
        if self.m != other.m:
            raise ValueError("Trees passed to query_ball_tree have different "
                             "dimensionality")
     
        n = self.n
        
        try:
        
            # allocate an array of std::vector<npy_intp>
            vvres = (<vector[np.intp_t] **> 
                PyMem_Malloc(n * sizeof(intvector_ptr_t)))
            if vvres == NULL:
                raise MemoryError()
                
            memset(<void*> vvres, 0, n * sizeof(intvector_ptr_t))      
            
            for i in range(n):
                vvres[i] = new vector[np.intp_t]()
        
            # query in C++
            # the GIL will be released in the C++ code
            query_ball_tree(
                <ckdtree*> self, <ckdtree*> other, r, p, eps, vvres)
                          
            # store the results in a list of lists                                        
            results = n * [None]
            for i in range(n):
                m = <np.intp_t> (vvres[i].size())
                if NPY_LIKELY(m > 0):
                    tmp = m * [None]
                    cur = npy_intp_vector_buf(vvres[i]) 
                    for j in range(m):
                        tmp[j] = cur[0]
                        cur += 1
                    results[i] = sorted(tmp)
                else:
                    results[i] = []    
                                  
        finally:
            if vvres != NULL:
                for i in range(n):
                    if vvres[i] != NULL:
                        del vvres[i]     
                PyMem_Free(vvres)

        return results

    # -----------
    # query_pairs
    # -----------
    
    def query_pairs(cKDTree self, np.float64_t r, np.float64_t p=2.,
                    np.float64_t eps=0, output_type='set'):
        """
        query_pairs(self, r, p=2., eps=0)

        Find all pairs of points whose distance is at most r.

        Parameters
        ----------
        r : positive float
            The maximum distance.
        p : float, optional
            Which Minkowski norm to use.  ``p`` has to meet the condition
            ``1 <= p <= infinity``.
        eps : float, optional
            Approximate search.  Branches of the tree are not explored
            if their nearest points are further than ``r/(1+eps)``, and
            branches are added in bulk if their furthest points are nearer
            than ``r * (1+eps)``.  `eps` has to be non-negative.
        output_type : string, optional
            Choose the output container, 'set' or 'ndarray'. Default: 'set'

        Returns
        -------
        results : set or ndarray
            Set of pairs ``(i,j)``, with ``i < j``, for which the corresponding
            positions are close. If output_type is 'ndarray', an ndarry is 
            returned instead of a set.

        """
                 
        cdef ordered_pairs results

        results = ordered_pairs()
        query_pairs(<ckdtree*> self, r, p, eps, results.buf)
        
        if output_type == 'set':
            return results.set()
        elif output_type == 'ndarray':
            return results.ndarray()
        else:
            raise ValueError("Invalid output type") 

    def _build_weights(cKDTree self, object weights):
        """
        _build_weights(weights)

        Compute weights of nodes from weights of data points. This will sum
        up the total weight per node. This function is used internally.

        Parameters
        ----------
        weights : array_like
            weights of data points; must be the same length as the data points.
            currently only scalar weights are supported. Therefore the weights
            array must be 1 dimensional.

        Returns
        -------
        node_weights : array_like
            total weight for each KD-Tree node.

        """
        cdef np.intp_t num_of_nodes
        cdef np.ndarray[np.float64_t, ndim=1, mode="c"] node_weights
        cdef np.ndarray[np.float64_t, ndim=1, mode="c"] proper_weights

        num_of_nodes = self.tree_buffer.size();
        node_weights = np.empty(num_of_nodes, dtype=np.float64)

        # FIXME: use templates to avoid the type conversion 
        proper_weights = np.ascontiguousarray(weights, dtype=np.float64)

        if len(proper_weights) != self.n:
            raise ValueError('Number of weights differ from the number of data points')

        build_weights(<ckdtree*> self, <np.float64_t*>np.PyArray_DATA(node_weights),
                            <np.float64_t*> np.PyArray_DATA(proper_weights))

        return node_weights

    # ---------------
    # count_neighbors
    # ---------------

    @cython.boundscheck(False)
    def count_neighbors(cKDTree self, cKDTree other, object r, np.float64_t p=2., 
                        object weights=None, int cumulative=True):
        """
        count_neighbors(self, other, r, p=2., weights=None, cumulative=True)

        Count how many nearby pairs can be formed. (pair-counting)

        Count the number of pairs (x1,x2) can be formed, with x1 drawn
        from self and x2 drawn from ``other``, and where
        ``distance(x1, x2, p) <= r``.

        Data points on self and other are optionally weighted by the ``weights``
        argument. (See below)

        The algorithm we implement here is based on [1]_. See notes for further discussion.

        Parameters
        ----------
        other : cKDTree instance
            The other tree to draw points from, can be the same tree as self.
        r : float or one-dimensional array of floats
            The radius to produce a count for. Multiple radii are searched with
            a single tree traversal. 
            If the count is non-cumulative(``cumulative=False``), ``r`` defines 
            the edges of the bins, and must be non-decreasing.
        p : float, optional 
            1<=p<=infinity. 
            Which Minkowski p-norm to use.
            Default 2.0.
        weights : tuple, array_like, or None, optional
            If None, the pair-counting is unweighted.
            If given as a tuple, weights[0] is the weights of points in ``self``, and
            weights[1] is the weights of points in ``other``; either can be None to 
            indicate the points are unweighted.
            If given as an array_like, weights is the weights of points in ``self``
            and ``other``. For this to make sense, ``self`` and ``other`` must be the
            same tree. If ``self`` and ``other`` are two different trees, a ``ValueError``
            is raised.
            Default: None
        cumulative : bool, optional
            Whether the returned counts are cumulative. When cumulative is set to ``False``
            the algorithm is optimized to work with a large number of bins (>10) specified
            by ``r``. When ``cumulative`` is set to True, the algorithm is optimized to work
            with a small number of ``r``. Default: True

        Returns
        -------
        result : scalar or 1-D array
            The number of pairs. For unweighted counts, the result is integer.
            For weighted counts, the result is float.
            If cumulative is False, ``result[i]`` contains the counts with
            ``(-inf if i == 0 else r[i-1]) < R <= r[i]``

        Notes
        -----
        Pair-counting is the basic operation used to calculate the two point
        correlation functions from a data set composed of position of objects.

        Two point correlation function measures the clustering of objects and
        is widely used in cosmology to quantify the large scale structure
        in our Universe, but it may be useful for data analysis in other fields
        where self-similar assembly of objects also occur.

        The Landy-Szalay estimator for the two point correlation function of
        ``D`` measures the clustering signal in ``D``. [2]_

        For example, given the position of two sets of objects,

        - objects ``D`` (data) contains the clustering signal, and

        - objects ``R`` (random) that contains no signal,

        .. math::

             \\xi(r) = \\frac{<D, D> - 2 f <D, R> + f^2<R, R>}{f^2<R, R>},

        where the brackets represents counting pairs between two data sets
        in a finite bin around ``r`` (distance), corresponding to setting
        `cumulative=False`, and ``f = float(len(D)) / float(len(R))`` is the
        ratio between number of objects from data and random.

        The algorithm implemented here is loosely based on the dual-tree
        algorithm described in [1]_. We switch between two different
        pair-cumulation scheme depending on the setting of ``cumulative``.
        The computing time of the method we use when for
        ``cumulative == False`` does not scale with the total number of bins.
        The algorithm for ``cumulative == True`` scales linearly with the
        number of bins, though it is slightly faster when only
        1 or 2 bins are used. [5]_.

        As an extension to the naive pair-counting,
        weighted pair-counting counts the product of weights instead
        of number of pairs.
        Weighted pair-counting is used to estimate marked correlation functions
        ([3]_, section 2.2),
        or to properly calculate the average of data per distance bin
        (e.g. [4]_, section 2.1 on redshift).

        .. [1] Gray and Moore,
               "N-body problems in statistical learning",
               Mining the sky, 2000,
               https://arxiv.org/abs/astro-ph/0012333

        .. [2] Landy and Szalay,
               "Bias and variance of angular correlation functions",
               The Astrophysical Journal, 1993,
               http://adsabs.harvard.edu/abs/1993ApJ...412...64L

        .. [3] Sheth, Connolly and Skibba,
               "Marked correlations in galaxy formation models",
               Arxiv e-print, 2005,
               https://arxiv.org/abs/astro-ph/0511773

        .. [4] Hawkins, et al.,
               "The 2dF Galaxy Redshift Survey: correlation functions,
               peculiar velocities and the matter density of the Universe",
               Monthly Notices of the Royal Astronomical Society, 2002,
               http://adsabs.harvard.edu/abs/2003MNRAS.346...78H

        .. [5] https://github.com/scipy/scipy/pull/5647#issuecomment-168474926

        """
        cdef: 
            int r_ndim
            np.intp_t n_queries, i
            np.ndarray[np.float64_t, ndim=1, mode="c"] real_r
            np.ndarray[np.float64_t, ndim=1, mode="c"] fresults
            np.ndarray[np.intp_t, ndim=1, mode="c"] iresults
            np.ndarray[np.float64_t, ndim=1, mode="c"] w1, w1n
            np.ndarray[np.float64_t, ndim=1, mode="c"] w2, w2n
            np.float64_t *w1p
            np.float64_t *w1np
            np.float64_t *w2p
            np.float64_t *w2np

        # Make sure trees are compatible
        if self.m != other.m:
            raise ValueError("Trees passed to count_neighbors have different "
                             "dimensionality")

        # Make a copy of r array to ensure it's contiguous and to modify it
        # below
        r_ndim = len(np.shape(r))
        if r_ndim > 1:
            raise ValueError("r must be either a single value or a "
                             "one-dimensional array of values")
        real_r = np.array(r, ndmin=1, dtype=np.float64, copy=True)
        if not cumulative:
            if (real_r[:-1] > real_r[1:]).any():
                raise ValueError("r must be non-decreasing for non-cumulative counting.");
        real_r, uind, inverse = np.unique(real_r, return_inverse=True, return_index=True)
        n_queries = real_r.shape[0]

        # Internally, we represent all distances as distance ** p
        if not ckdtree_isinf(p):
            for i in range(n_queries):
                if not ckdtree_isinf(real_r[i]):
                    real_r[i] = real_r[i] ** p

        if weights is None:
            self_weights = other_weights = None
        elif isinstance(weights, tuple):
            self_weights, other_weights = weights
        else:
            self_weights = other_weights = weights
            if other is not self:
                raise ValueError("Two different trees are used. Specify weights for both in a tuple.")

        if self_weights is None and other_weights is None:
            int_result = True
            # unweighted, use the integer arithmetics
            results = np.zeros(n_queries + 1, dtype=np.intp)

            iresults = results
            count_neighbors_unweighted(<ckdtree*> self, <ckdtree*> other, n_queries,
                            &real_r[0], &iresults[0], p, cumulative)

        else:
            int_result = False
            # weighted / half weighted, use the floating point arithmetics
            if self_weights is not None:
                w1 = np.ascontiguousarray(self_weights, dtype=np.float64)
                w1n = self._build_weights(w1)
                w1p = <np.float64_t*> np.PyArray_DATA(w1)
                w1np = <np.float64_t*> np.PyArray_DATA(w1n)
            else:
                w1p = NULL
                w1np = NULL
            if other_weights is not None:
                w2 = np.ascontiguousarray(other_weights, dtype=np.float64)
                w2n = other._build_weights(w2)
                w2p = <np.float64_t*> np.PyArray_DATA(w2)
                w2np = <np.float64_t*> np.PyArray_DATA(w2n)
            else:
                w2p = NULL
                w2np = NULL

            results = np.zeros(n_queries + 1, dtype=np.float64)
            fresults = results
            count_neighbors_weighted(<ckdtree*> self, <ckdtree*> other,
                                    w1p, w2p, w1np, w2np,
                                    n_queries,
                                    &real_r[0], &fresults[0], p, cumulative)

        results2 = np.zeros(inverse.shape, results.dtype)
        if cumulative:
            # copy out the results (taking care of duplication and sorting)
            results2[...] = results[inverse]
        else:
            # keep the identical ones zero
            # this could have been done in a more readable way.
            results2[uind] = results[inverse][uind]
        results = results2

        if r_ndim == 0:
            if int_result and results[0] <= <np.intp_t> LONG_MAX:
                return int(results[0])
            else:
                return results[0]
        else:
            return results
    
    # ----------------------
    # sparse_distance_matrix
    # ----------------------
    
    def sparse_distance_matrix(cKDTree self, cKDTree other,
                               np.float64_t max_distance,
                               np.float64_t p=2.,
                               output_type='dok_matrix'):
        """
        sparse_distance_matrix(self, other, max_distance, p=2.)

        Compute a sparse distance matrix

        Computes a distance matrix between two cKDTrees, leaving as zero
        any distance greater than max_distance.

        Parameters
        ----------
        other : cKDTree

        max_distance : positive float
        
        p : float, 1<=p<=infinity
            Which Minkowski p-norm to use. 
        
        output_type : string, optional
            Which container to use for output data. Options: 'dok_matrix',
            'coo_matrix', 'dict', or 'ndarray'. Default: 'dok_matrix'.

        Returns
        -------
        result : dok_matrix, coo_matrix, dict or ndarray
            Sparse matrix representing the results in "dictionary of keys" 
            format. If a dict is returned the keys are (i,j) tuples of indices.
            If output_type is 'ndarray' a record array with fields 'i', 'j',
            and 'k' is returned,
        """
        
        cdef coo_entries res

        # Make sure trees are compatible
        if self.m != other.m:
            raise ValueError("Trees passed to sparse_distance_matrix have "
                             "different dimensionality")                                      
        # do the query
        res = coo_entries()
        sparse_distance_matrix(
                <ckdtree*> self, <ckdtree*> other, p, max_distance, res.buf)
                
        if output_type == 'dict':
            return res.dict()
        elif output_type == 'ndarray':
            return res.ndarray()
        elif output_type == 'coo_matrix':
            return res.coo_matrix(self.n, other.n)            
        elif output_type == 'dok_matrix':
            return res.dok_matrix(self.n, other.n)
        else:
            raise ValueError('Invalid output type')


    # ----------------------
    # pickle
    # ----------------------    

        
    def __reduce__(self):
        return (new_object, (cKDTree,), self.__getstate__())

    def __getstate__(cKDTree self):
        cdef object state
        cdef object tree = pickle_tree_buffer(self.tree_buffer)
        state = (tree, self.data.copy(), self.n, self.m, self.leafsize,
                      self.maxes, self.mins, self.indices.copy(), 
                      self.boxsize, self.boxsize_data)
        return state
            
    def __setstate__(cKDTree self, state):
        cdef object tree
        self.tree_buffer = new vector[ckdtreenode]()
        
        # unpack the state
        (tree, self.data, self.n, self.m, self.leafsize, 
            self.maxes, self.mins, self.indices, self.boxsize, self.boxsize_data) = state

        # set raw pointers
        self._pre_init()
        
        # copy kd-tree buffer 
        unpickle_tree_buffer(self.tree_buffer, tree)    
        
        # set up the tree structure pointers
        self._post_init()
        
        # make the tree viewable from Python
        self.tree = cKDTreeNode()
        self.tree._node = self.ctree
        self.tree._data = self.data
        self.tree._indices = self.indices
        self.tree.level = 0
        self.tree._setup()