/***********************************************************************
 * The code in this file is modified from the original code in the
 * https://github.com/jlblancoc/nanoflann repository.
 * 
 * Software License Agreement (BSD License)
 *
 * Copyright 2008-2009  Marius Muja (mariusm@cs.ubc.ca). All rights reserved.
 * Copyright 2008-2009  David G. Lowe (lowe@cs.ubc.ca). All rights reserved.
 * Copyright 2011-2025  Jose Luis Blanco (joseluisblancoc@gmail.com).
 *   All rights reserved.
 *
 * THE BSD LICENSE
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *************************************************************************/

#include "cuda_kdtree.cuh"
// #define THRUST_DEBUG 1
#include <cuda.h>
#include <thrust/gather.h>
#include <thrust/copy.h>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/copy.h>
#include <thrust/scan.h>
#include <thrust/count.h>

#include <limits>

__device__ __forceinline__
float infinity()
{
    return __int_as_float(0x7f800000);
}

#ifndef INFINITY
#define INFINITY infinity()
#endif

namespace tinyflann
{
    namespace cuda {

        template <class T>
        __device__ __host__ void swap( T& x, T& y ) {
            T t=x; x=y; y=t;
        }

        template <class GreaterThan, class RandomAccessIterator, class RandomAccessIterator2>
        __host__ __device__ void
        sift_down( RandomAccessIterator key, RandomAccessIterator2 value, size_t begin, size_t length, GreaterThan c = GreaterThan() )
        {

            while( 2*begin+1 < length ) {
                size_t left = 2*begin+1;
                size_t right = 2*begin+2;
                size_t largest=begin;
                if((left < length)&& c(key[left], key[largest]) ) largest=left;

                if((right < length)&& c(key[right], key[largest]) ) largest=right;

                if( largest != begin ) {
                    cuda::swap( key[begin], key[largest] );
                    cuda::swap( value[begin], value[largest] );
                    begin=largest;
                }
                else return;
            }
        }

//! creates a max-heap in the array beginning at begin of length "length"
//! if c is a less-than comparator, it will create a min-heap
        template <class GreaterThan, class RandomAccessIterator, class RandomAccessIterator2>
        __host__ __device__ void
        make_heap( RandomAccessIterator key,  RandomAccessIterator2 value, size_t length, GreaterThan c = GreaterThan() )
        {
            int i=length/2-1;
            while( i>=0 ) {
                sift_down( key, value, i, length, c );
                i--;
            }
        }

        template< typename DistanceType >
        struct GreaterThan
        {
            __device__
            bool operator()(DistanceType a, DistanceType b)
            {
                return a>b;
            }
        };

//! result set for the 1nn search. Doesn't do any global memory accesses on its own,
        template<typename DistanceType>
        struct SingleResultSet {
            int bestIndex;
            DistanceType bestDist;
            const DistanceType epsError;

            __device__ __host__ SingleResultSet(DistanceType eps) : bestIndex(-1), bestDist(INFINITY), epsError(eps) {}

            __device__
            inline float
            worstDist() {
                return bestDist;
            }

            __device__
            inline void
            insert(int index, DistanceType dist) {
                if (dist <= bestDist) {
                    bestIndex = index;
                    bestDist = dist;
                }
            }

            DistanceType *resultDist;
            int *resultIndex;

            __device__
            inline void
            setResultLocation(DistanceType *dists, int *index, int thread, int stride) {
                resultDist = dists + thread * stride;
                resultIndex = index + thread * stride;
                if (stride != 1) {
                    for (int i = 1; i < stride; i++) {
                        resultDist[i] = INFINITY;
                        resultIndex[i] = -1;
                    }
                }
            }

            __device__
            inline void
            finish() {
                resultDist[0] = bestDist;
                resultIndex[0] = bestIndex;
            }
        };

        template<typename DistanceType, bool useHeap>
        struct KnnResultSet {
            int foundNeighbors;
            DistanceType largestHeapDist;
            int maxDistIndex;
            const int k;
            const bool sorted;
            const DistanceType epsError;


            __device__ __host__
            KnnResultSet(int knn, bool sortResults, DistanceType eps) : foundNeighbors(0), largestHeapDist(INFINITY),
                                                                        k(knn), sorted(sortResults), epsError(eps) {}

            __device__
            inline DistanceType
            worstDist() {
                return largestHeapDist;
            }

            __device__
            inline void
            insert(int index, DistanceType dist) {
                if (foundNeighbors < k) {
                    resultDist[foundNeighbors] = dist;
                    resultIndex[foundNeighbors] = index;
                    if (foundNeighbors == k - 1) {
                        if (useHeap) {
                            make_heap(resultDist, resultIndex, k, GreaterThan<DistanceType>());
                            largestHeapDist = resultDist[0];
                        } else {
                            findLargestDistIndex();
                        }

                    }
                    foundNeighbors++;
                } else if (dist < largestHeapDist) {
                    if (useHeap) {
                        resultDist[0] = dist;
                        resultIndex[0] = index;
                        sift_down(resultDist, resultIndex, 0, k, GreaterThan<DistanceType>());
                        largestHeapDist = resultDist[0];
                    } else {
                        resultDist[maxDistIndex] = dist;
                        resultIndex[maxDistIndex] = index;
                        findLargestDistIndex();
                    }

                }
            }

            __device__
            void
            findLargestDistIndex() {
                largestHeapDist = resultDist[0];
                maxDistIndex = 0;
                for (int i = 1; i < k; i++)
                    if (resultDist[i] > largestHeapDist) {
                        maxDistIndex = i;
                        largestHeapDist = resultDist[i];
                    }
            }

            float *resultDist;
            int *resultIndex;

            __device__ inline void setResultLocation(DistanceType *dists, int *index, int thread, int stride) {
                resultDist = dists + stride * thread;
                resultIndex = index + stride * thread;
                for (int i = 0; i < stride; i++) {
                    resultDist[i] = INFINITY;
                    resultIndex[i] = -1;
                }
            }

            __host__ __device__ inline void finish() {
                if (sorted) {
                    if (!useHeap) make_heap(resultDist, resultIndex, k, GreaterThan<DistanceType>());
                    for (int i = k - 1; i > 0; i--) {
                        cuda::swap(resultDist[0], resultDist[i]);
                        cuda::swap(resultIndex[0], resultIndex[i]);
                        sift_down(resultDist, resultIndex, 0, i, GreaterThan<DistanceType>());
                    }
                }
            }
        };

        template<typename DistanceType>
        struct CountingRadiusResultSet {
            int count_;
            DistanceType radius_sq_;
            int max_neighbors_;

            __device__ __host__
            CountingRadiusResultSet(DistanceType radius, int max_neighbors) : count_(0), radius_sq_(radius),
                                                                              max_neighbors_(max_neighbors) {}

            __device__
            inline DistanceType
            worstDist() {
                return radius_sq_;
            }

            __device__
            inline void
            insert(int index, float dist) {
                if (dist < radius_sq_) {
                    count_++;
                }
            }

            int *resultIndex;

            __device__
            inline void
            setResultLocation(DistanceType * /*dists*/, int *count, int thread, int stride) {
                resultIndex = count + thread * stride;
            }

            __device__
            inline void
            finish() {
                if ((max_neighbors_ <= 0) || (count_ <= max_neighbors_)) resultIndex[0] = count_;
                else resultIndex[0] = max_neighbors_;
            }
        };

// Difference to RadiusKnnResultSet: Works like KnnResultSet, doesn't pack the results densely (as the RadiusResultSet does)
        template<typename DistanceType, bool useHeap>
        struct KnnRadiusResultSet {
            int foundNeighbors;
            DistanceType largestHeapDist;
            int maxDistIndex;
            const int k;
            const bool sorted;
            const DistanceType epsError;
            const DistanceType radius_sq;


            __device__ __host__
            KnnRadiusResultSet(int knn, bool sortResults, DistanceType eps, DistanceType radius) : foundNeighbors(0),
                                                                                                   largestHeapDist(
                                                                                                           radius),
                                                                                                   k(knn),
                                                                                                   sorted(sortResults),
                                                                                                   epsError(eps),
                                                                                                   radius_sq(radius) {}

            __device__
            inline DistanceType
            worstDist() {
                return largestHeapDist;
            }

            __device__
            inline void
            insert(int index, DistanceType dist) {
                if (dist < largestHeapDist) {
                    if (foundNeighbors < k) {
                        resultDist[foundNeighbors] = dist;
                        resultIndex[foundNeighbors] = index;
                        if (foundNeighbors == k - 1) {
                            if (useHeap) {
                                make_heap(resultDist, resultIndex, k, GreaterThan<DistanceType>());
                                largestHeapDist = resultDist[0];
                            } else {
                                findLargestDistIndex();
                            }
                        }
                        foundNeighbors++;
                    } else { //if( dist < largestHeapDist )
                        if (useHeap) {
                            resultDist[0] = dist;
                            resultIndex[0] = index;
                            sift_down(resultDist, resultIndex, 0, k, GreaterThan<DistanceType>());
                            largestHeapDist = resultDist[0];
                        } else {
                            resultDist[maxDistIndex] = dist;
                            resultIndex[maxDistIndex] = index;
                            findLargestDistIndex();
                        }
                    }
                }
            }

            __device__
            void
            findLargestDistIndex() {
                largestHeapDist = resultDist[0];
                maxDistIndex = 0;
                for (int i = 1; i < k; i++)
                    if (resultDist[i] > largestHeapDist) {
                        maxDistIndex = i;
                        largestHeapDist = resultDist[i];
                    }
            }

            DistanceType *resultDist;
            int *resultIndex;

            __device__ inline void setResultLocation(DistanceType *dists, int *index, int thread, int stride) {
                resultDist = dists + stride * thread;
                resultIndex = index + stride * thread;
                for (int i = 0; i < stride; i++) {
                    resultDist[i] = INFINITY;
                    resultIndex[i] = -1;
                }
            }

            __device__ inline void finish() {
                if (sorted) {
                    if (!useHeap) make_heap(resultDist, resultIndex, k, GreaterThan<DistanceType>());
                    for (int i = k - 1; i > 0; i--) {
                        cuda::swap(resultDist[0], resultDist[i]);
                        cuda::swap(resultIndex[0], resultIndex[i]);
                        sift_down(resultDist, resultIndex, 0, i, GreaterThan<DistanceType>());
                    }
                }
            }
        };
    }
}


namespace tinyflann {

    namespace cuda {
        namespace kd_tree_builder_detail {
//! normal node: contains the split dimension and value
//! leaf node: left == index of first points, right==index of last point +1
            struct SplitInfo {
                union {
                    struct {
                        // begin of child nodes
                        int left;
                        // end of child nodes
                        int right;
                    };
                    struct {
                        int split_dim;
                        float split_val;
                    };
                };

            };

            struct IsEven {
                typedef int result_type;

                __device__
                int operator()(int i) {
                    return (i & 1) == 0;
                }
            };

            struct SecondElementIsEven {
                __host__ __device__
                bool operator()(const thrust::tuple<int, int> &i) {
                    return (thrust::get<1>(i) & 1) == 0;
                }
            };

//! just for convenience: access a float4 by an index in [0,1,2]
//! (casting it to a float* and accessing it by the index is way slower...)
            __host__ __device__
            float get_value_by_index(const float4 &f, int i) {
                switch (i) {
                    case 0:
                        return f.x;
                    case 1:
                        return f.y;
                    default:
                        return f.z;
                }

            }

//! mark a point as belonging to the left or right child of its current parent
//! called after parents are split
            struct MovePointsToChildNodes {
                MovePointsToChildNodes(int *child1, SplitInfo *splits, float *x, float *y, float *z, int *ox, int *oy,
                                       int *oz, int *lrx, int *lry, int *lrz)
                        : child1_(child1), splits_(splits), x_(x), y_(y), z_(z), ox_(ox), oy_(oy), oz_(oz), lrx_(lrx),
                          lry_(lry), lrz_(lrz) {}

                //  int dim;
                //  float threshold;
                int *child1_;
                SplitInfo *splits_;

                // coordinate values
                float *x_, *y_, *z_;
                // owner indices -> which node does the point belong to?
                int *ox_, *oy_, *oz_;
                // temp info: will be set to 1 of a point is moved to the right child node, 0 otherwise
                // (used later in the scan op to separate the points of the children into continuous ranges)
                int *lrx_, *lry_, *lrz_;

                __device__
                void operator()(const thrust::tuple<int, int, int, int> &data) {
                    int index = thrust::get<0>(data);
                    int owner = ox_[index]; // before a split, all points at the same position in the index array have the same owner
                    int point_ind1 = thrust::get<1>(data);
                    int point_ind2 = thrust::get<2>(data);
                    int point_ind3 = thrust::get<3>(data);
                    int leftChild = child1_[owner];
                    int split_dim;
                    float dim_val1, dim_val2, dim_val3;
                    SplitInfo split;
                    lrx_[index] = 0;
                    lry_[index] = 0;
                    lrz_[index] = 0;
                    // this element already belongs to a leaf node -> everything alright, no need to change anything
                    if (leftChild == -1) {
                        return;
                    }
                    // otherwise: load split data, and assign this index to the new owner
                    split = splits_[owner];
                    split_dim = split.split_dim;
                    switch (split_dim) {
                        case 0:
                            dim_val1 = x_[point_ind1];
                            dim_val2 = x_[point_ind2];
                            dim_val3 = x_[point_ind3];
                            break;
                        case 1:
                            dim_val1 = y_[point_ind1];
                            dim_val2 = y_[point_ind2];
                            dim_val3 = y_[point_ind3];
                            break;
                        default:
                            dim_val1 = z_[point_ind1];
                            dim_val2 = z_[point_ind2];
                            dim_val3 = z_[point_ind3];
                            break;

                    }


                    int r1 = leftChild + (dim_val1 > split.split_val);
                    ox_[index] = r1;
                    int r2 = leftChild + (dim_val2 > split.split_val);
                    oy_[index] = r2;
                    oz_[index] = leftChild + (dim_val3 > split.split_val);

                    lrx_[index] = (dim_val1 > split.split_val);
                    lry_[index] = (dim_val2 > split.split_val);
                    lrz_[index] = (dim_val3 > split.split_val);
                    //                      return thrust::make_tuple( r1, r2, leftChild+(dim_val > split.split_val) );
                }
            };

//! used to update the left/right pointers and aabb infos after the node splits
            struct SetLeftAndRightAndAABB {
                int maxPoints;
                int nElements;

                SplitInfo *nodes;
                int *counts;
                int *labels;
                float4 *aabbMin;
                float4 *aabbMax;
                const float *x, *y, *z;
                const int *ix, *iy, *iz;

                __host__ __device__
                void operator()(int i) {
                    int index = labels[i];
                    int right;
                    int left = counts[i];
                    nodes[index].left = left;
                    if (i < nElements - 1) {
                        right = counts[i + 1];
                    } else { // index==nNodes
                        right = maxPoints;
                    }
                    nodes[index].right = right;
                    aabbMin[index].x = x[ix[left]];
                    aabbMin[index].y = y[iy[left]];
                    aabbMin[index].z = z[iz[left]];
                    aabbMax[index].x = x[ix[right - 1]];
                    aabbMax[index].y = y[iy[right - 1]];
                    aabbMax[index].z = z[iz[right - 1]];
                }
            };


//! - decide whether a node has to be split
//! if yes:
//! - allocate child nodes
//! - set split axis as axis of maximum aabb length
            struct SplitNodes {
                int maxPointsPerNode;
                int *node_count;
                int *nodes_allocated;
                int *out_of_space;
                int *child1_;
                int *parent_;
                SplitInfo *splits;

                __device__
                void operator()(
                        thrust::tuple<int &, int &, SplitInfo &, float4 &, float4 &, int> node) // float4: aabbMin, aabbMax
                {
                    int &parent = thrust::get<0>(node);
                    int &child1 = thrust::get<1>(node);
                    SplitInfo &s = thrust::get<2>(node);
                    const float4 &aabbMin = thrust::get<3>(node);
                    const float4 &aabbMax = thrust::get<4>(node);
                    int my_index = thrust::get<5>(node);
                    bool split_node = false;
                    // first, each thread block counts the number of nodes that it needs to allocate...
                    __shared__ int block_nodes_to_allocate;
                    if (threadIdx.x == 0) block_nodes_to_allocate = 0;
                    __syncthreads();

                    // don't split if all points are equal
                    // (could lead to an infinite loop, and doesn't make any sense anyway)
                    bool all_points_in_node_are_equal =
                            aabbMin.x == aabbMax.x && aabbMin.y == aabbMax.y && aabbMin.z == aabbMax.z;

                    int offset_to_global = 0;

                    // maybe this could be replaced with a reduction...
                    if ((child1 == -1) && (s.right - s.left > maxPointsPerNode) &&
                        !all_points_in_node_are_equal) { // leaf node
                        split_node = true;
                        offset_to_global = atomicAdd(&block_nodes_to_allocate, 2);
                    }

                    __syncthreads();
                    __shared__ int block_left;
                    __shared__ bool enough_space;
                    // ... then the first thread tries to allocate this many nodes...
                    if (threadIdx.x == 0) {
                        block_left = atomicAdd(node_count, block_nodes_to_allocate);
                        enough_space = block_left + block_nodes_to_allocate < *nodes_allocated;
                        // if it doesn't succeed, no nodes will be created by this block
                        if (!enough_space) {
                            atomicAdd(node_count, -block_nodes_to_allocate);
                            *out_of_space = 1;
                        }
                    }

                    __syncthreads();
                    // this thread needs to split it's node && there was enough space for all the nodes
                    // in this block.
                    //(The whole "allocate-per-block-thing" is much faster than letting each element allocate
                    // its space on its own, because shared memory atomics are A LOT faster than
                    // global mem atomics!)
                    if (split_node && enough_space) {
                        int left = block_left + offset_to_global;

                        splits[left].left = s.left;
                        splits[left].right = s.right;
                        splits[left + 1].left = 0;
                        splits[left + 1].right = 0;

                        // split axis/position: middle of longest aabb extent
                        float4 aabbDim = aabbMax - aabbMin;
                        int maxDim = 0;
                        float maxDimLength = aabbDim.x;
                        float4 splitVal = (aabbMax + aabbMin);
                        splitVal *= 0.5f;
                        for (int i = 1; i <= 2; i++) {
                            float val = get_value_by_index(aabbDim, i);
                            if (val > maxDimLength) {
                                maxDim = i;
                                maxDimLength = val;
                            }
                        }
                        s.split_dim = maxDim;
                        s.split_val = get_value_by_index(splitVal, maxDim);

                        child1_[my_index] = left;
                        splits[my_index] = s;

                        parent_[left] = my_index;
                        parent_[left + 1] = my_index;
                        child1_[left] = -1;
                        child1_[left + 1] = -1;
                    }
                }
            };


//! computes the scatter target address for the split operation, see Sengupta,Harris,Zhang,Owen: Scan Primitives for GPU Computing
//! in my use case, this is about 2x as fast as thrust::partition
            struct set_addr3 {
                const int *val_, *f_;

                int npoints_;

                __device__
                int operator()(int id) {
                    int nf = f_[npoints_ - 1] + (val_[npoints_ - 1]);
                    int f = f_[id];
                    int t = id - f + nf;
                    return val_[id] ? f : t;
                }
            };

//! converts a float4 point (xyz) to a tuple of three float vals (used to separate the
//! float4 input buffer into three arrays in the beginning of the tree build)
            struct pointxyz_to_px_py_pz {
                __device__
                thrust::tuple<float, float, float> operator()(const float4 &val) {
                    return thrust::make_tuple(val.x, val.y, val.z);
                }
            };
        } // namespace kd_tree_builder_detail

    } // namespace cuda


    class CudaKdTreeBuilder {
    public:
        CudaKdTreeBuilder(const thrust::device_vector<float4> &points, int max_leaf_size)
                : /*out_of_space_(1,0),node_count_(1,1),*/ max_leaf_size_(max_leaf_size) {
            points_ = &points;
            int prealloc = points.size() / max_leaf_size_ * 16;
            allocation_info_.resize(3);
            allocation_info_[NodeCount] = 1;
            allocation_info_[NodesAllocated] = prealloc;
            allocation_info_[OutOfSpace] = 0;

            child1_ = new thrust::device_vector<int>(prealloc, -1);
            parent_ = new thrust::device_vector<int>(prealloc, -1);
            cuda::kd_tree_builder_detail::SplitInfo s;
            s.left = 0;
            s.right = 0;
            splits_ = new thrust::device_vector<cuda::kd_tree_builder_detail::SplitInfo>(prealloc, s);
            s.right = points.size();
            (*splits_)[0] = s;

            aabb_min_ = new thrust::device_vector<float4>(prealloc);
            aabb_max_ = new thrust::device_vector<float4>(prealloc);

            index_x_ = new thrust::device_vector<int>(points_->size());
            index_y_ = new thrust::device_vector<int>(points_->size());
            index_z_ = new thrust::device_vector<int>(points_->size());

            owners_x_ = new thrust::device_vector<int>(points_->size(), 0);
            owners_y_ = new thrust::device_vector<int>(points_->size(), 0);
            owners_z_ = new thrust::device_vector<int>(points_->size(), 0);

            leftright_x_ = new thrust::device_vector<int>(points_->size(), 0);
            leftright_y_ = new thrust::device_vector<int>(points_->size(), 0);
            leftright_z_ = new thrust::device_vector<int>(points_->size(), 0);

            tmp_index_ = new thrust::device_vector<int>(points_->size());
            tmp_owners_ = new thrust::device_vector<int>(points_->size());
            tmp_misc_ = new thrust::device_vector<int>(points_->size());

            points_x_ = new thrust::device_vector<float>(points_->size());
            points_y_ = new thrust::device_vector<float>(points_->size());
            points_z_ = new thrust::device_vector<float>(points_->size());
            delete_node_info_ = false;
        }

        ~CudaKdTreeBuilder() {
            if (delete_node_info_) {
                delete child1_;
                delete parent_;
                delete splits_;
                delete aabb_min_;
                delete aabb_max_;
                delete index_x_;
            }

            delete index_y_;
            delete index_z_;
            delete owners_x_;
            delete owners_y_;
            delete owners_z_;
            delete points_x_;
            delete points_y_;
            delete points_z_;
            delete leftright_x_;
            delete leftright_y_;
            delete leftright_z_;
            delete tmp_index_;
            delete tmp_owners_;
            delete tmp_misc_;
        }

        //! build the tree
        //! general idea:
        //! - build sorted lists of the points in x y and z order (to be able to compute tight AABBs in O(1) )
        //! - while( nodes to split exist )
        //!    - split non-child nodes along longest axis if the number of points is > max_points_per_node
        //!    - for each point: determine whether it is in a node that was split. If yes, mark it as belonging to the left or right child node of its current parent node
        //!    - reorder the points so that the points of a single node are continuous in the node array
        //!    - update the left/right pointers and AABBs of all nodes
        void buildTree() {
            //              std::cout<<"buildTree()"<<std::endl;
            //              sleep(1);
            //              Util::Timer stepTimer;
            thrust::transform(points_->begin(), points_->end(), thrust::make_zip_iterator(
                                      thrust::make_tuple(points_x_->begin(), points_y_->begin(), points_z_->begin())),
                              cuda::kd_tree_builder_detail::pointxyz_to_px_py_pz());

            thrust::counting_iterator<int> it(0);
            thrust::copy(it, it + points_->size(), index_x_->begin());

            thrust::copy(index_x_->begin(), index_x_->end(), index_y_->begin());
            thrust::copy(index_x_->begin(), index_x_->end(), index_z_->begin());

            thrust::device_vector<float> tmpv(points_->size());

            // create sorted index list -> can be used to compute AABBs in O(1)
            thrust::copy(points_x_->begin(), points_x_->end(), tmpv.begin());
            thrust::sort_by_key(tmpv.begin(), tmpv.end(), index_x_->begin());
            thrust::copy(points_y_->begin(), points_y_->end(), tmpv.begin());
            thrust::sort_by_key(tmpv.begin(), tmpv.end(), index_y_->begin());
            thrust::copy(points_z_->begin(), points_z_->end(), tmpv.begin());
            thrust::sort_by_key(tmpv.begin(), tmpv.end(), index_z_->begin());


            (*aabb_min_)[0] = make_float4((*points_x_)[(*index_x_)[0]], (*points_y_)[(*index_y_)[0]],
                                          (*points_z_)[(*index_z_)[0]], 0);

            (*aabb_max_)[0] = make_float4((*points_x_)[(*index_x_)[points_->size() - 1]],
                                          (*points_y_)[(*index_y_)[points_->size() - 1]],
                                          (*points_z_)[(*index_z_)[points_->size() - 1]], 0);
#ifdef PRINT_DEBUG_TIMING
            cudaDeviceSynchronize();
        std::cout<<" initial stuff:"<<stepTimer.elapsed()<<std::endl;
        stepTimer.restart();
#endif
            int last_node_count = 0;
            for (int i = 0;; i++) {
                cuda::kd_tree_builder_detail::SplitNodes sn;

                sn.maxPointsPerNode = max_leaf_size_;
                sn.node_count = thrust::raw_pointer_cast(&allocation_info_[NodeCount]);
                sn.nodes_allocated = thrust::raw_pointer_cast(&allocation_info_[NodesAllocated]);
                sn.out_of_space = thrust::raw_pointer_cast(&allocation_info_[OutOfSpace]);
                sn.child1_ = thrust::raw_pointer_cast(&(*child1_)[0]);
                sn.parent_ = thrust::raw_pointer_cast(&(*parent_)[0]);
                sn.splits = thrust::raw_pointer_cast(&(*splits_)[0]);

                thrust::counting_iterator<int> cit(0);
                thrust::for_each(thrust::make_zip_iterator(
                                         thrust::make_tuple(parent_->begin(), child1_->begin(), splits_->begin(), aabb_min_->begin(),
                                                            aabb_max_->begin(), cit)),
                                 thrust::make_zip_iterator(thrust::make_tuple(parent_->begin() + last_node_count,
                                                                              child1_->begin() + last_node_count,
                                                                              splits_->begin() + last_node_count,
                                                                              aabb_min_->begin() + last_node_count,
                                                                              aabb_max_->begin() + last_node_count,
                                                                              cit + last_node_count)),
                                 sn);
                // copy allocation info to host
                thrust::host_vector<int> alloc_info = allocation_info_;

                if (last_node_count == alloc_info[NodeCount]) { // no more nodes were split -> done
                    break;
                }
                last_node_count = alloc_info[NodeCount];

                // a node was un-splittable due to a lack of space
                if (alloc_info[OutOfSpace] == 1) {
                    resize_node_vectors(alloc_info[NodesAllocated] * 2);
                    alloc_info[OutOfSpace] = 0;
                    alloc_info[NodesAllocated] *= 2;
                    allocation_info_ = alloc_info;
                }
#ifdef PRINT_DEBUG_TIMING
                cudaDeviceSynchronize();
            std::cout<<" node split:"<<stepTimer.elapsed()<<std::endl;
            stepTimer.restart();
#endif

                // foreach point: point was in node that was split?move it to child (leaf) node : do nothing
                cuda::kd_tree_builder_detail::MovePointsToChildNodes sno(thrust::raw_pointer_cast(&(*child1_)[0]),
                                                                         thrust::raw_pointer_cast(&(*splits_)[0]),
                                                                         thrust::raw_pointer_cast(&(*points_x_)[0]),
                                                                         thrust::raw_pointer_cast(&(*points_y_)[0]),
                                                                         thrust::raw_pointer_cast(&(*points_z_)[0]),
                                                                         thrust::raw_pointer_cast(&(*owners_x_)[0]),
                                                                         thrust::raw_pointer_cast(&(*owners_y_)[0]),
                                                                         thrust::raw_pointer_cast(&(*owners_z_)[0]),
                                                                         thrust::raw_pointer_cast(&(*leftright_x_)[0]),
                                                                         thrust::raw_pointer_cast(&(*leftright_y_)[0]),
                                                                         thrust::raw_pointer_cast(&(*leftright_z_)[0])
                );
                thrust::counting_iterator<int> ci0(0);
                thrust::for_each(thrust::make_zip_iterator(
                                         thrust::make_tuple(ci0, index_x_->begin(), index_y_->begin(), index_z_->begin())),
                                 thrust::make_zip_iterator(
                                         thrust::make_tuple(ci0 + points_->size(), index_x_->end(), index_y_->end(),
                                                            index_z_->end())), sno);

#ifdef PRINT_DEBUG_TIMING
                cudaDeviceSynchronize();
            std::cout<<" set new owners:"<<stepTimer.elapsed()<<std::endl;
            stepTimer.restart();
#endif

                // move points around so that each leaf node's points are continuous
                separate_left_and_right_children(*index_x_, *owners_x_, *tmp_index_, *tmp_owners_, *leftright_x_);
                std::swap(tmp_index_, index_x_);
                std::swap(tmp_owners_, owners_x_);
                separate_left_and_right_children(*index_y_, *owners_y_, *tmp_index_, *tmp_owners_, *leftright_y_,
                                                 false);
                std::swap(tmp_index_, index_y_);
                separate_left_and_right_children(*index_z_, *owners_z_, *tmp_index_, *tmp_owners_, *leftright_z_,
                                                 false);
                std::swap(tmp_index_, index_z_);

#ifdef PRINT_DEBUG_TIMING
                cudaDeviceSynchronize();
            std::cout<<" split:"<<stepTimer.elapsed()<<std::endl;
            stepTimer.restart();
#endif
                // calculate new AABB etc
                update_leftright_and_aabb(*points_x_, *points_y_, *points_z_, *index_x_, *index_y_, *index_z_,
                                          *owners_x_, *splits_, *aabb_min_, *aabb_max_);
#ifdef PRINT_DEBUG_TIMING
                cudaDeviceSynchronize();
            std::cout<<" update_leftright_and_aabb:"<<stepTimer.elapsed()<<std::endl;
            stepTimer.restart();
            print_vector(node_count_);
#endif

            }
        }

        template<class Distance>
        friend
        class KDTreeCuda3dIndex;

    protected:


        //! takes the partitioned nodes, and sets the left-/right info of leaf nodes, as well as the AABBs
        void
        update_leftright_and_aabb(const thrust::device_vector<float> &x, const thrust::device_vector<float> &y,
                                  const thrust::device_vector<float> &z,
                                  const thrust::device_vector<int> &ix, const thrust::device_vector<int> &iy,
                                  const thrust::device_vector<int> &iz,
                                  const thrust::device_vector<int> &owners,
                                  thrust::device_vector<cuda::kd_tree_builder_detail::SplitInfo> &splits,
                                  thrust::device_vector<float4> &aabbMin, thrust::device_vector<float4> &aabbMax) {
            thrust::device_vector<int> *labelsUnique = tmp_owners_;
            thrust::device_vector<int> *countsUnique = tmp_index_;
            // assume: points of each node are continuous in the array

            // find which nodes are here, and where each node's points begin and end
            int unique_labels =
                    thrust::unique_by_key_copy(owners.begin(), owners.end(), thrust::counting_iterator<int>(0),
                                               labelsUnique->begin(), countsUnique->begin()).first -
                    labelsUnique->begin();

            // update the info
            cuda::kd_tree_builder_detail::SetLeftAndRightAndAABB s;
            s.maxPoints = x.size();
            s.nElements = unique_labels;
            s.nodes = thrust::raw_pointer_cast(&(splits[0]));
            s.counts = thrust::raw_pointer_cast(&((*countsUnique)[0]));
            s.labels = thrust::raw_pointer_cast(&((*labelsUnique)[0]));
            s.x = thrust::raw_pointer_cast(&x[0]);
            s.y = thrust::raw_pointer_cast(&y[0]);
            s.z = thrust::raw_pointer_cast(&z[0]);
            s.ix = thrust::raw_pointer_cast(&ix[0]);
            s.iy = thrust::raw_pointer_cast(&iy[0]);
            s.iz = thrust::raw_pointer_cast(&iz[0]);
            s.aabbMin = thrust::raw_pointer_cast(&aabbMin[0]);
            s.aabbMax = thrust::raw_pointer_cast(&aabbMax[0]);

            thrust::counting_iterator<int> it(0);
            thrust::for_each(it, it + unique_labels, s);
        }

        //! Separates the left and right children of each node into continuous parts of the array.
        //! More specifically, it seperates children with even and odd node indices because nodes are always
        //! allocated in pairs -> child1==child2+1 -> child1 even and child2 odd, or vice-versa.
        //! Since the split operation is stable, this results in continuous partitions
        //! for all the single nodes.
        //! (basically the split primitive according to sengupta et al)
        //! about twice as fast as thrust::partition
        void separate_left_and_right_children(thrust::device_vector<int> &key_in, thrust::device_vector<int> &val_in,
                                              thrust::device_vector<int> &key_out, thrust::device_vector<int> &val_out,
                                              thrust::device_vector<int> &left_right_marks,
                                              bool scatter_val_out = true) {
            thrust::device_vector<int> *f_tmp = &val_out;
            thrust::device_vector<int> *addr_tmp = tmp_misc_;

            thrust::exclusive_scan( /*thrust::make_transform_iterator(*/
                    left_right_marks.begin() /*,cuda::kd_tree_builder_detail::IsEven*/
                    /*())*/, /*thrust::make_transform_iterator(*/
                    left_right_marks.end() /*,cuda::kd_tree_builder_detail::IsEven*/
                    /*())*/, f_tmp->begin());
            cuda::kd_tree_builder_detail::set_addr3 sa;
            sa.val_ = thrust::raw_pointer_cast(&left_right_marks[0]);
            sa.f_ = thrust::raw_pointer_cast(&(*f_tmp)[0]);
            sa.npoints_ = key_in.size();
            thrust::counting_iterator<int> it(0);
            thrust::transform(it, it + val_in.size(), addr_tmp->begin(), sa);

            thrust::scatter(key_in.begin(), key_in.end(), addr_tmp->begin(), key_out.begin());
            if (scatter_val_out) thrust::scatter(val_in.begin(), val_in.end(), addr_tmp->begin(), val_out.begin());
        }

        //! allocates additional space in all the node-related vectors.
        //! new_size elements will be added to all vectors.
        void resize_node_vectors(size_t new_size) {
            size_t add = new_size - child1_->size();
            child1_->insert(child1_->end(), add, -1);
            parent_->insert(parent_->end(), add, -1);
            cuda::kd_tree_builder_detail::SplitInfo s;
            s.left = 0;
            s.right = 0;
            splits_->insert(splits_->end(), add, s);
            float4 f;
            aabb_min_->insert(aabb_min_->end(), add, f);
            aabb_max_->insert(aabb_max_->end(), add, f);
        }


        const thrust::device_vector<float4> *points_;

        // tree data, those are stored per-node

        //! left child of each node. (right child==left child + 1, due to the alloc mechanism)
        //! child1_[node]==-1 if node is a leaf node
        thrust::device_vector<int> *child1_;
        //! parent node of each node
        thrust::device_vector<int> *parent_;
        //! split info (dim/value or left/right pointers)
        thrust::device_vector<cuda::kd_tree_builder_detail::SplitInfo> *splits_;
        //! min aabb value of each node
        thrust::device_vector<float4> *aabb_min_;
        //! max aabb value of each node
        thrust::device_vector<float4> *aabb_max_;

        enum AllocationInfo {
            NodeCount = 0,
            NodesAllocated = 1,
            OutOfSpace = 2
        };
        // those were put into a single vector of 3 elements so that only one mem transfer will be needed for all three of them
        //  thrust::device_vector<int> out_of_space_;
        //  thrust::device_vector<int> node_count_;
        //  thrust::device_vector<int> nodes_allocated_;
        thrust::device_vector<int> allocation_info_;

        int max_leaf_size_;

        // coordinate values of the points
        thrust::device_vector<float> *points_x_, *points_y_, *points_z_;
        // indices
        thrust::device_vector<int> *index_x_, *index_y_, *index_z_;
        // owner node
        thrust::device_vector<int> *owners_x_, *owners_y_, *owners_z_;
        // contains info about whether a point was partitioned to the left or right child after a split
        thrust::device_vector<int> *leftright_x_, *leftright_y_, *leftright_z_;
        thrust::device_vector<int> *tmp_index_, *tmp_owners_, *tmp_misc_;
        bool delete_node_info_;
    };


    namespace KdTreeCudaPrivate {
        template<typename GPUResultSet, typename Distance>
        __device__ void searchNeighbors(const cuda::kd_tree_builder_detail::SplitInfo *splits,
                                        const int *child1,
                                        const int *parent,
                                        const float4 *aabbLow,
                                        const float4 *aabbHigh, const float4 *elements, const float4 &q, GPUResultSet &result,
                                        const Distance &distance = Distance()) {

            bool backtrack = false;
            int lastNode = -1;
            int current = 0;

            cuda::kd_tree_builder_detail::SplitInfo split;
            while (true) {
                if (current == -1) break;
                split = splits[current];

                float diff1;
                if (split.split_dim == 0) diff1 = q.x - split.split_val;
                else if (split.split_dim == 1) diff1 = q.y - split.split_val;
                else if (split.split_dim == 2) diff1 = q.z - split.split_val;

                // children are next to each other: leftChild+1 == rightChild
                int leftChild = child1[current];
                int bestChild = leftChild;
                int otherChild = leftChild;

                if ((diff1) < 0) {
                    otherChild++;
                } else {
                    bestChild++;
                }

                if (!backtrack) {
                    /* If this is a leaf node, then do check and return. */
                    if (leftChild == -1) {
                        for (int i = split.left; i < split.right; ++i) {
                            float dist = distance.dist(elements[i], q);
                            result.insert(i, dist);
                        }
                        backtrack = true;
                        lastNode = current;
                        current = parent[current];
                    } else { // go to closer child node
                        lastNode = current;
                        current = bestChild;
                    }
                } else { // continue moving back up the tree or visit far node?
                    // minimum possible distance between query point and a point inside the AABB
                    float mindistsq = 0;
                    float4 aabbMin = aabbLow[otherChild];
                    float4 aabbMax = aabbHigh[otherChild];

                    if (q.x < aabbMin.x) mindistsq += distance.axisDist(q.x, aabbMin.x);
                    else if (q.x > aabbMax.x) mindistsq += distance.axisDist(q.x, aabbMax.x);
                    if (q.y < aabbMin.y) mindistsq += distance.axisDist(q.y, aabbMin.y);
                    else if (q.y > aabbMax.y) mindistsq += distance.axisDist(q.y, aabbMax.y);
                    if (q.z < aabbMin.z) mindistsq += distance.axisDist(q.z, aabbMin.z);
                    else if (q.z > aabbMax.z) mindistsq += distance.axisDist(q.z, aabbMax.z);

                    //  the far node was NOT the last node (== not visited yet) AND there could be a closer point in it
                    if ((lastNode == bestChild) && (mindistsq <= result.worstDist())) {
                        lastNode = current;
                        current = otherChild;
                        backtrack = false;
                    } else {
                        lastNode = current;
                        current = parent[current];
                    }
                }

            }
        }


        template<typename GPUResultSet, typename Distance>
        __global__
        void nearestKernel(const cuda::kd_tree_builder_detail::SplitInfo *splits,
                           const int *child1,
                           const int *parent,
                           const float4 *aabbMin,
                           const float4 *aabbMax, const float4 *elements, const float *query, int stride,
                           int resultStride, int *resultIndex, float *resultDist, int querysize, GPUResultSet result,
                           Distance dist = Distance()) {
            typedef float DistanceType;
            typedef float ElementType;
            //                  typedef DistanceType float;
            size_t tid = blockDim.x * blockIdx.x + threadIdx.x;

            if (tid >= querysize) return;

            float4 q = make_float4(query[tid * stride], query[tid * stride + 1], query[tid * stride + 2], 0);

            result.setResultLocation(resultDist, resultIndex, tid, resultStride);

            searchNeighbors(splits, child1, parent, aabbMin, aabbMax, elements, q, result, dist);

            result.finish();
        }

    }

//! contains some pointers that use cuda data types and that cannot be easily
//! forward-declared.
//! basically it contains all GPU buffers
    template<typename Distance>
    struct KDTreeCuda3dIndex<Distance>::GpuHelper {
        thrust::device_vector<cuda::kd_tree_builder_detail::SplitInfo> *gpu_splits_;
        thrust::device_vector<int> *gpu_parent_;
        thrust::device_vector<int> *gpu_child1_;
        thrust::device_vector<float4> *gpu_aabb_min_;
        thrust::device_vector<float4> *gpu_aabb_max_;
        thrust::device_vector<float4> *gpu_points_;
        thrust::device_vector<int> *gpu_vind_;

        GpuHelper() : gpu_splits_(0), gpu_parent_(0), gpu_child1_(0), gpu_aabb_min_(0), gpu_aabb_max_(0),
                      gpu_points_(0), gpu_vind_(0) {
        }

        ~GpuHelper() {
            delete gpu_splits_;
            gpu_splits_ = 0;
            delete gpu_parent_;
            gpu_parent_ = 0;
            delete gpu_child1_;
            gpu_child1_ = 0;
            delete gpu_aabb_max_;
            gpu_aabb_max_ = 0;
            delete gpu_aabb_min_;
            gpu_aabb_min_ = 0;
            delete gpu_vind_;
            gpu_vind_ = 0;

            delete gpu_points_;
            gpu_points_ = 0;
        }
    };

//! thrust transform functor
//! transforms indices in the internal data set back to the original indices
    struct map_indices {
        const int *v_;

        map_indices(const int *v) : v_(v) {}

        __host__ __device__
        float operator()(const int &i) const {
            if (i >= 0) return v_[i];
            else return i;
        }
    };

//! implementation of L2 distance for the CUDA kernels
    struct CudaL2 {

        static float __host__ __device__ axisDist(float a, float b) {
            return (a - b) * (a - b);
        }

        static float __host__ __device__ dist(float4 a, float4 b) {
            float4 diff = a - b;
            return dot(diff, diff);
        }
    };

//! implementation of L1 distance for the CUDA kernels
//! NOT TESTED!
    struct CudaL1 {

        static float __host__ __device__ axisDist(float a, float b) {
            return fabs(a - b);
        }

        static float __host__ __device__ dist(float4 a, float4 b) {
            return fabs(a.x - b.x) + fabs(a.y - b.y) + (a.z - b.z) + (a.w - b.w);
        }
    };


    template<typename Distance>
    void KDTreeCuda3dIndex<Distance>::knnSearch(const float* queries, size_t n_query, int istride, int* indices,
                                                float* dists, size_t knn, const SearchParams& params) const {
        int ostride = knn;

        int threadsPerBlock = 128;
        int blocksPerGrid = (n_query + threadsPerBlock - 1) / threadsPerBlock;

        float epsError = 1 + params.eps;
        bool sorted = params.sorted;
        bool use_heap = params.use_heap;

        Distance distance;
        if (knn == 1) {
            KdTreeCudaPrivate::nearestKernel<<<blocksPerGrid, threadsPerBlock >> > (thrust::raw_pointer_cast(
                    &((*gpu_helper_->gpu_splits_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_child1_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_parent_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_aabb_min_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_aabb_max_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_points_)[0])),
                    queries,
                    istride,
                    ostride,
                    indices,
                    dists,
                    n_query, cuda::SingleResultSet<float>(epsError), distance);
        } else {
            if (use_heap) {
                KdTreeCudaPrivate::nearestKernel<<<blocksPerGrid, threadsPerBlock >> > (thrust::raw_pointer_cast(
                        &((*gpu_helper_->gpu_splits_)[0])),
                        thrust::raw_pointer_cast(&((*gpu_helper_->gpu_child1_)[0])),
                        thrust::raw_pointer_cast(&((*gpu_helper_->gpu_parent_)[0])),
                        thrust::raw_pointer_cast(&((*gpu_helper_->gpu_aabb_min_)[0])),
                        thrust::raw_pointer_cast(&((*gpu_helper_->gpu_aabb_max_)[0])),
                        thrust::raw_pointer_cast(&((*gpu_helper_->gpu_points_)[0])),
                        queries,
                        istride,
                        ostride,
                        indices,
                        dists,
                        n_query, cuda::KnnResultSet<float, true>(knn, sorted, epsError), distance);
            } else {
                KdTreeCudaPrivate::nearestKernel<<<blocksPerGrid, threadsPerBlock >> > (thrust::raw_pointer_cast(
                        &((*gpu_helper_->gpu_splits_)[0])),
                        thrust::raw_pointer_cast(&((*gpu_helper_->gpu_child1_)[0])),
                        thrust::raw_pointer_cast(&((*gpu_helper_->gpu_parent_)[0])),
                        thrust::raw_pointer_cast(&((*gpu_helper_->gpu_aabb_min_)[0])),
                        thrust::raw_pointer_cast(&((*gpu_helper_->gpu_aabb_max_)[0])),
                        thrust::raw_pointer_cast(&((*gpu_helper_->gpu_points_)[0])),
                        queries,
                        istride,
                        ostride,
                        indices,
                        dists,
                        n_query, cuda::KnnResultSet<float, false>(knn, sorted, epsError), distance
                );
            }
        }
        thrust::device_ptr<int> id = thrust::device_pointer_cast(indices);
        thrust::transform(id, id + knn * n_query, id,
                          map_indices(thrust::raw_pointer_cast(&((*gpu_helper_->gpu_vind_))[0])));
    }
//! used in the radius search to count the total number of neighbors
    struct isNotMinusOne {
        __host__ __device__
        bool operator()(int i) {
            return i != -1;
        }
    };
    template<typename Distance>
    int KDTreeCuda3dIndex<Distance>::radiusSearch(const float* queries, size_t n_query, int n_query_stride, int* indices,
                                                  float* dists, float radius, const SearchParams& params) const {
        int max_neighbors = params.max_neighbors;
        assert(max_neighbors > 0);
        bool sorted = params.sorted;
        float epsError = 1 + params.eps;
        bool use_heap = params.use_heap;
        int istride = n_query_stride;
        int ostride = max_neighbors;
        Distance distance;
        int threadsPerBlock = 128;
        int blocksPerGrid = (n_query + threadsPerBlock - 1) / threadsPerBlock;
        if (use_heap) {
            KdTreeCudaPrivate::nearestKernel<<<blocksPerGrid, threadsPerBlock >> > (thrust::raw_pointer_cast(
                    &((*gpu_helper_->gpu_splits_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_child1_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_parent_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_aabb_min_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_aabb_max_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_points_)[0])),
                    queries,
                    istride,
                    ostride,
                    indices,
                    dists,
                    n_query, cuda::KnnRadiusResultSet<float, true>(max_neighbors, sorted, epsError,
                                                                   radius), distance);
        } else {
            KdTreeCudaPrivate::nearestKernel<<<blocksPerGrid, threadsPerBlock >> > (thrust::raw_pointer_cast(
                    &((*gpu_helper_->gpu_splits_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_child1_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_parent_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_aabb_min_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_aabb_max_)[0])),
                    thrust::raw_pointer_cast(&((*gpu_helper_->gpu_points_)[0])),
                    queries,
                    istride,
                    ostride,
                    indices,
                    dists,
                    n_query, cuda::KnnRadiusResultSet<float, false>(max_neighbors, sorted, epsError,
                                                                    radius), distance);
        }
        thrust::device_ptr<int> id = thrust::device_pointer_cast(indices);
        thrust::transform(id, id + max_neighbors * n_query, id,
                          map_indices(thrust::raw_pointer_cast(&((*gpu_helper_->gpu_vind_))[0])));
        return thrust::count_if(id, id + max_neighbors * n_query, isNotMinusOne());
    }
    template<typename Distance>
    void KDTreeCuda3dIndex<Distance>::uploadTreeToGpu() {
        // just make sure that no weird alignment stuff is going on...
        // shouldn't, but who knows
        // (I would make this a (boost) static assertion, but so far flann seems to avoid boost
        //  assert( sizeof( KdTreeCudaPrivate::GpuNode)==sizeof( Node ) );
        delete gpu_helper_;
        gpu_helper_ = new GpuHelper;
        gpu_helper_->gpu_points_ = new thrust::device_vector<float4>(size_);
        thrust::device_vector<float4> tmp(size_);
        thrust::copy(thrust::device_pointer_cast((float4 *) dataset_),
                     thrust::device_pointer_cast((float4 *) (dataset_)) + size_, tmp.begin());
        CudaKdTreeBuilder builder(tmp, leaf_max_size_);
        builder.buildTree();
        gpu_helper_->gpu_splits_ = builder.splits_;
        gpu_helper_->gpu_aabb_min_ = builder.aabb_min_;
        gpu_helper_->gpu_aabb_max_ = builder.aabb_max_;
        gpu_helper_->gpu_child1_ = builder.child1_;
        gpu_helper_->gpu_parent_ = builder.parent_;
        gpu_helper_->gpu_vind_ = builder.index_x_;
        thrust::gather(builder.index_x_->begin(), builder.index_x_->end(), tmp.begin(), gpu_helper_->gpu_points_->begin());
    }
    template<typename Distance>
    void KDTreeCuda3dIndex<Distance>::clearGpuBuffers() {
        delete gpu_helper_;
        gpu_helper_ = 0;
    }
    // instantiations for L1 and L2.
    template void KDTreeCuda3dIndex<CudaL1>::uploadTreeToGpu();
    template void KDTreeCuda3dIndex<CudaL2>::uploadTreeToGpu();
    template void KDTreeCuda3dIndex<CudaL1>::clearGpuBuffers();
    template void KDTreeCuda3dIndex<CudaL2>::clearGpuBuffers();
    template struct KDTreeCuda3dIndex<CudaL1 >::GpuHelper;
    template struct KDTreeCuda3dIndex<CudaL2 >::GpuHelper;
    template void KDTreeCuda3dIndex<CudaL1>::knnSearch(const float* queries, size_t n_query, int n_query_stride, int* indices, float* dists, size_t knn, const SearchParams& params) const;
    template void KDTreeCuda3dIndex<CudaL2>::knnSearch(const float* queries, size_t n_query, int n_query_stride, int* indices, float* dists, size_t knn, const SearchParams& params) const;
    template int KDTreeCuda3dIndex<CudaL1>::radiusSearch(const float* queries, size_t n_query, int n_query_stride, int* indices, float* dists, float radius, const SearchParams& params) const;
    template int KDTreeCuda3dIndex<CudaL2>::radiusSearch(const float* queries, size_t n_query, int n_query_stride, int* indices, float* dists, float radius, const SearchParams& params) const;
};