ceres-solver and colmap

2b5a2b6 almost 4 years ago

24.2 kB

/***********************************************************************
 * 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.
 *
 * 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.
 *************************************************************************/

#ifndef FLANN_KDTREE_INDEX_H_
#define FLANN_KDTREE_INDEX_H_

#include <algorithm>
#include <map>
#include <cassert>
#include <cstring>
#include <stdarg.h>
#include <cmath>
#include <random>

#include "FLANN/general.h"
#include "FLANN/algorithms/nn_index.h"
#include "FLANN/util/dynamic_bitset.h"
#include "FLANN/util/matrix.h"
#include "FLANN/util/result_set.h"
#include "FLANN/util/heap.h"
#include "FLANN/util/allocator.h"
#include "FLANN/util/random.h"
#include "FLANN/util/saving.h"


namespace flann
{

struct KDTreeIndexParams : public IndexParams
{
    KDTreeIndexParams(int trees = 4)
    {
        (*this)["algorithm"] = FLANN_INDEX_KDTREE;
        (*this)["trees"] = trees;
    }
};


/**
 * Randomized kd-tree index
 *
 * Contains the k-d trees and other information for indexing a set of points
 * for nearest-neighbor matching.
 */
template <typename Distance>
class KDTreeIndex : public NNIndex<Distance>
{
public:
    typedef typename Distance::ElementType ElementType;
    typedef typename Distance::ResultType DistanceType;

    typedef NNIndex<Distance> BaseClass;

    typedef bool needs_kdtree_distance;


    /**
     * KDTree constructor
     *
     * Params:
     *          inputData = dataset with the input features
     *          params = parameters passed to the kdtree algorithm
     */
    KDTreeIndex(const IndexParams& params = KDTreeIndexParams(), Distance d = Distance() ) :
    	BaseClass(params, d), mean_(NULL), var_(NULL)
    {
        trees_ = get_param(index_params_,"trees",4);
    }


    /**
     * KDTree constructor
     *
     * Params:
     *          inputData = dataset with the input features
     *          params = parameters passed to the kdtree algorithm
     */
    KDTreeIndex(const Matrix<ElementType>& dataset, const IndexParams& params = KDTreeIndexParams(),
                Distance d = Distance() ) : BaseClass(params,d ), mean_(NULL), var_(NULL)
    {
        trees_ = get_param(index_params_,"trees",4);

        setDataset(dataset);
    }

    KDTreeIndex(const KDTreeIndex& other) : BaseClass(other),
    		trees_(other.trees_)
    {
        tree_roots_.resize(other.tree_roots_.size());
        for (size_t i=0;i<tree_roots_.size();++i) {
        	copyTree(tree_roots_[i], other.tree_roots_[i]);
        }
    }

    KDTreeIndex& operator=(KDTreeIndex other)
    {
    	this->swap(other);
    	return *this;
    }

    /**
     * Standard destructor
     */
    virtual ~KDTreeIndex()
    {
    	freeIndex();
    }

    BaseClass* clone() const
    {
    	return new KDTreeIndex(*this);
    }

    using BaseClass::buildIndex;

    void addPoints(const Matrix<ElementType>& points, float rebuild_threshold = 2)
    {
        assert(points.cols==veclen_);

        size_t old_size = size_;
        extendDataset(points);

        if (rebuild_threshold>1 && size_at_build_*rebuild_threshold<size_) {
            buildIndex();
        }
        else {
            for (size_t i=old_size;i<size_;++i) {
                for (int j = 0; j < trees_; j++) {
                    addPointToTree(tree_roots_[j], i);
                }
            }
        }
    }

    flann_algorithm_t getType() const
    {
        return FLANN_INDEX_KDTREE;
    }


    template<typename Archive>
    void serialize(Archive& ar)
    {
    	ar.setObject(this);

    	ar & *static_cast<NNIndex<Distance>*>(this);

    	ar & trees_;

    	if (Archive::is_loading::value) {
    		tree_roots_.resize(trees_);
    	}
    	for (size_t i=0;i<tree_roots_.size();++i) {
    		if (Archive::is_loading::value) {
    			tree_roots_[i] = new(pool_) Node();
    		}
    		ar & *tree_roots_[i];
    	}

    	if (Archive::is_loading::value) {
            index_params_["algorithm"] = getType();
            index_params_["trees"] = trees_;
    	}
    }


    void saveIndex(FILE* stream)
    {
    	serialization::SaveArchive sa(stream);
    	sa & *this;
    }


    void loadIndex(FILE* stream)
    {
    	freeIndex();
    	serialization::LoadArchive la(stream);
    	la & *this;
    }

    /**
     * Computes the inde memory usage
     * Returns: memory used by the index
     */
    int usedMemory() const
    {
        return int(pool_.usedMemory+pool_.wastedMemory+size_*sizeof(int));  // pool memory and vind array memory
    }

    /**
     * Find set of nearest neighbors to vec. Their indices are stored inside
     * the result object.
     *
     * Params:
     *     result = the result object in which the indices of the nearest-neighbors are stored
     *     vec = the vector for which to search the nearest neighbors
     *     maxCheck = the maximum number of restarts (in a best-bin-first manner)
     */
    void findNeighbors(ResultSet<DistanceType>& result, const ElementType* vec, const SearchParams& searchParams) const
    {
        int maxChecks = searchParams.checks;
        float epsError = 1+searchParams.eps;

        if (maxChecks==FLANN_CHECKS_UNLIMITED) {
        	if (removed_) {
        		getExactNeighbors<true>(result, vec, epsError);
        	}
        	else {
        		getExactNeighbors<false>(result, vec, epsError);
        	}
        }
        else {
        	if (removed_) {
        		getNeighbors<true>(result, vec, maxChecks, epsError);
        	}
        	else {
        		getNeighbors<false>(result, vec, maxChecks, epsError);
        	}
        }
    }

protected:

    /**
     * Builds the index
     */
    void buildIndexImpl()
    {
        // Create a permutable array of indices to the input vectors.
    	std::vector<int> ind(size_);
        for (size_t i = 0; i < size_; ++i) {
            ind[i] = int(i);
        }

        mean_ = new DistanceType[veclen_];
        var_ = new DistanceType[veclen_];

        std::default_random_engine generator;

        tree_roots_.resize(trees_);
        /* Construct the randomized trees. */
        for (int i = 0; i < trees_; i++) {
            /* Randomize the order of vectors to allow for unbiased sampling. */
            std::shuffle(ind.begin(), ind.end(), generator);
            tree_roots_[i] = divideTree(&ind[0], int(size_) );
        }
        delete[] mean_;
        delete[] var_;
    }

    void freeIndex()
    {
    	for (size_t i=0;i<tree_roots_.size();++i) {
    		// using placement new, so call destructor explicitly
    		if (tree_roots_[i]!=NULL) tree_roots_[i]->~Node();
    	}
    	pool_.free();
    }


private:

    /*--------------------- Internal Data Structures --------------------------*/
    struct Node
    {
    	/**
         * Dimension used for subdivision.
         */
        int divfeat;
        /**
         * The values used for subdivision.
         */
        DistanceType divval;
        /**
         * Point data
         */
        ElementType* point;
		/**
		* The child nodes.
		*/
		Node* child1, *child2;
		Node(){
			child1 = NULL;
			child2 = NULL;
		}
		~Node() {
			if (child1 != NULL) { child1->~Node(); child1 = NULL; }

			if (child2 != NULL) { child2->~Node(); child2 = NULL; }
		}

    private:
    	template<typename Archive>
    	void serialize(Archive& ar)
    	{
    		typedef KDTreeIndex<Distance> Index;
    		Index* obj = static_cast<Index*>(ar.getObject());

    		ar & divfeat;
    		ar & divval;

    		bool leaf_node = false;
    		if (Archive::is_saving::value) {
    			leaf_node = ((child1==NULL) && (child2==NULL));
    		}
    		ar & leaf_node;

    		if (leaf_node) {
    			if (Archive::is_loading::value) {
    				point = obj->points_[divfeat];
    			}
    		}

    		if (!leaf_node) {
				if (Archive::is_loading::value) {
					child1 = new(obj->pool_) Node();
					child2 = new(obj->pool_) Node();
				}
    			ar & *child1;
    			ar & *child2;
    		}
    	}
    	friend struct serialization::access;
    };
    typedef Node* NodePtr;
    typedef BranchStruct<NodePtr, DistanceType> BranchSt;
    typedef BranchSt* Branch;


    void copyTree(NodePtr& dst, const NodePtr& src)
    {
    	dst = new(pool_) Node();
    	dst->divfeat = src->divfeat;
    	dst->divval = src->divval;
    	if (src->child1==NULL && src->child2==NULL) {
    		dst->point = points_[dst->divfeat];
    		dst->child1 = NULL;
    		dst->child2 = NULL;
    	}
    	else {
    		copyTree(dst->child1, src->child1);
    		copyTree(dst->child2, src->child2);
    	}
    }

    /**
     * Create a tree node that subdivides the list of vecs from vind[first]
     * to vind[last].  The routine is called recursively on each sublist.
     * Place a pointer to this new tree node in the location pTree.
     *
     * Params: pTree = the new node to create
     *                  first = index of the first vector
     *                  last = index of the last vector
     */
    NodePtr divideTree(int* ind, int count)
    {
        NodePtr node = new(pool_) Node(); // allocate memory

        /* If too few exemplars remain, then make this a leaf node. */
        if (count == 1) {
            node->child1 = node->child2 = NULL;    /* Mark as leaf node. */
            node->divfeat = *ind;    /* Store index of this vec. */
            node->point = points_[*ind];
        }
        else {
            int idx;
            int cutfeat;
            DistanceType cutval;
            meanSplit(ind, count, idx, cutfeat, cutval);

            node->divfeat = cutfeat;
            node->divval = cutval;
            node->child1 = divideTree(ind, idx);
            node->child2 = divideTree(ind+idx, count-idx);
        }

        return node;
    }


    /**
     * Choose which feature to use in order to subdivide this set of vectors.
     * Make a random choice among those with the highest variance, and use
     * its variance as the threshold value.
     */
    void meanSplit(int* ind, int count, int& index, int& cutfeat, DistanceType& cutval)
    {
        memset(mean_,0,veclen_*sizeof(DistanceType));
        memset(var_,0,veclen_*sizeof(DistanceType));

        /* Compute mean values.  Only the first SAMPLE_MEAN values need to be
            sampled to get a good estimate.
         */
        int cnt = std::min((int)SAMPLE_MEAN+1, count);
        for (int j = 0; j < cnt; ++j) {
            ElementType* v = points_[ind[j]];
            for (size_t k=0; k<veclen_; ++k) {
                mean_[k] += v[k];
            }
        }
        DistanceType div_factor = DistanceType(1)/cnt;
        for (size_t k=0; k<veclen_; ++k) {
            mean_[k] *= div_factor;
        }

        /* Compute variances (no need to divide by count). */
        for (int j = 0; j < cnt; ++j) {
            ElementType* v = points_[ind[j]];
            for (size_t k=0; k<veclen_; ++k) {
                DistanceType dist = v[k] - mean_[k];
                var_[k] += dist * dist;
            }
        }
        /* Select one of the highest variance indices at random. */
        cutfeat = selectDivision(var_);
        cutval = mean_[cutfeat];

        int lim1, lim2;
        planeSplit(ind, count, cutfeat, cutval, lim1, lim2);

        if (lim1>count/2) index = lim1;
        else if (lim2<count/2) index = lim2;
        else index = count/2;

        /* If either list is empty, it means that all remaining features
         * are identical. Split in the middle to maintain a balanced tree.
         */
        if ((lim1==count)||(lim2==0)) index = count/2;
    }


    /**
     * Select the top RAND_DIM largest values from v and return the index of
     * one of these selected at random.
     */
    int selectDivision(DistanceType* v)
    {
        int num = 0;
        size_t topind[RAND_DIM];

        /* Create a list of the indices of the top RAND_DIM values. */
        for (size_t i = 0; i < veclen_; ++i) {
            if ((num < RAND_DIM)||(v[i] > v[topind[num-1]])) {
                /* Put this element at end of topind. */
                if (num < RAND_DIM) {
                    topind[num++] = i;            /* Add to list. */
                }
                else {
                    topind[num-1] = i;         /* Replace last element. */
                }
                /* Bubble end value down to right location by repeated swapping. */
                int j = num - 1;
                while (j > 0  &&  v[topind[j]] > v[topind[j-1]]) {
                    std::swap(topind[j], topind[j-1]);
                    --j;
                }
            }
        }
        /* Select a random integer in range [0,num-1], and return that index. */
        int rnd = rand_int(num);
        return (int)topind[rnd];
    }


    /**
     *  Subdivide the list of points by a plane perpendicular on axe corresponding
     *  to the 'cutfeat' dimension at 'cutval' position.
     *
     *  On return:
     *  dataset[ind[0..lim1-1]][cutfeat]<cutval
     *  dataset[ind[lim1..lim2-1]][cutfeat]==cutval
     *  dataset[ind[lim2..count]][cutfeat]>cutval
     */
    void planeSplit(int* ind, int count, int cutfeat, DistanceType cutval, int& lim1, int& lim2)
    {
        /* Move vector indices for left subtree to front of list. */
        int left = 0;
        int right = count-1;
        for (;; ) {
            while (left<=right && points_[ind[left]][cutfeat]<cutval) ++left;
            while (left<=right && points_[ind[right]][cutfeat]>=cutval) --right;
            if (left>right) break;
            std::swap(ind[left], ind[right]); ++left; --right;
        }
        lim1 = left;
        right = count-1;
        for (;; ) {
            while (left<=right && points_[ind[left]][cutfeat]<=cutval) ++left;
            while (left<=right && points_[ind[right]][cutfeat]>cutval) --right;
            if (left>right) break;
            std::swap(ind[left], ind[right]); ++left; --right;
        }
        lim2 = left;
    }

    /**
     * Performs an exact nearest neighbor search. The exact search performs a full
     * traversal of the tree.
     */
    template<bool with_removed>
    void getExactNeighbors(ResultSet<DistanceType>& result, const ElementType* vec, float epsError) const
    {
        //		checkID -= 1;  /* Set a different unique ID for each search. */

        if (trees_ > 1) {
            fprintf(stderr,"It doesn't make any sense to use more than one tree for exact search");
        }
        if (trees_>0) {
            searchLevelExact<with_removed>(result, vec, tree_roots_[0], 0.0, epsError);
        }
    }

    /**
     * Performs the approximate nearest-neighbor search. The search is approximate
     * because the tree traversal is abandoned after a given number of descends in
     * the tree.
     */
    template<bool with_removed>
    void getNeighbors(ResultSet<DistanceType>& result, const ElementType* vec, int maxCheck, float epsError) const
    {
        int i;
        BranchSt branch;

        int checkCount = 0;
        Heap<BranchSt>* heap = new Heap<BranchSt>((int)size_);
        DynamicBitset checked(size_);

        /* Search once through each tree down to root. */
        for (i = 0; i < trees_; ++i) {
            searchLevel<with_removed>(result, vec, tree_roots_[i], 0, checkCount, maxCheck, epsError, heap, checked);
        }

        /* Keep searching other branches from heap until finished. */
        while ( heap->popMin(branch) && (checkCount < maxCheck || !result.full() )) {
            searchLevel<with_removed>(result, vec, branch.node, branch.mindist, checkCount, maxCheck, epsError, heap, checked);
        }

        delete heap;

    }

    /**
     *  Search starting from a given node of the tree.  Based on any mismatches at
     *  higher levels, all exemplars below this level must have a distance of
     *  at least "mindistsq".
     */
    template<bool with_removed>
    void searchLevel(ResultSet<DistanceType>& result_set, const ElementType* vec, NodePtr node, DistanceType mindist, int& checkCount, int maxCheck,
                     float epsError, Heap<BranchSt>* heap, DynamicBitset& checked) const
    {
        if (result_set.worstDist()<mindist) {
            //			printf("Ignoring branch, too far\n");
            return;
        }

        /* If this is a leaf node, then do check and return. */
        if ((node->child1 == NULL)&&(node->child2 == NULL)) {
            int index = node->divfeat;
            if (with_removed) {
            	if (removed_points_.test(index)) return;
            }
            /*  Do not check same node more than once when searching multiple trees. */
            if ( checked.test(index) || ((checkCount>=maxCheck)&& result_set.full()) ) return;
            checked.set(index);
            checkCount++;

            DistanceType dist = distance_(node->point, vec, veclen_);
            result_set.addPoint(dist,index);
            return;
        }

        /* Which child branch should be taken first? */
        ElementType val = vec[node->divfeat];
        DistanceType diff = val - node->divval;
        NodePtr bestChild = (diff < 0) ? node->child1 : node->child2;
        NodePtr otherChild = (diff < 0) ? node->child2 : node->child1;

        /* Create a branch record for the branch not taken.  Add distance
            of this feature boundary (we don't attempt to correct for any
            use of this feature in a parent node, which is unlikely to
            happen and would have only a small effect).  Don't bother
            adding more branches to heap after halfway point, as cost of
            adding exceeds their value.
         */

        DistanceType new_distsq = mindist + distance_.accum_dist(val, node->divval, node->divfeat);
        //		if (2 * checkCount < maxCheck  ||  !result.full()) {
        if ((new_distsq*epsError < result_set.worstDist())||  !result_set.full()) {
            heap->insert( BranchSt(otherChild, new_distsq) );
        }

        /* Call recursively to search next level down. */
        searchLevel<with_removed>(result_set, vec, bestChild, mindist, checkCount, maxCheck, epsError, heap, checked);
    }

    /**
     * Performs an exact search in the tree starting from a node.
     */
    template<bool with_removed>
    void searchLevelExact(ResultSet<DistanceType>& result_set, const ElementType* vec, const NodePtr node, DistanceType mindist, const float epsError) const
    {
        /* If this is a leaf node, then do check and return. */
        if ((node->child1 == NULL)&&(node->child2 == NULL)) {
            int index = node->divfeat;
            if (with_removed) {
            	if (removed_points_.test(index)) return; // ignore removed points
            }
            DistanceType dist = distance_(node->point, vec, veclen_);
            result_set.addPoint(dist,index);

            return;
        }

        /* Which child branch should be taken first? */
        ElementType val = vec[node->divfeat];
        DistanceType diff = val - node->divval;
        NodePtr bestChild = (diff < 0) ? node->child1 : node->child2;
        NodePtr otherChild = (diff < 0) ? node->child2 : node->child1;

        /* Create a branch record for the branch not taken.  Add distance
            of this feature boundary (we don't attempt to correct for any
            use of this feature in a parent node, which is unlikely to
            happen and would have only a small effect).  Don't bother
            adding more branches to heap after halfway point, as cost of
            adding exceeds their value.
         */

        DistanceType new_distsq = mindist + distance_.accum_dist(val, node->divval, node->divfeat);

        /* Call recursively to search next level down. */
        searchLevelExact<with_removed>(result_set, vec, bestChild, mindist, epsError);

        if (mindist*epsError<=result_set.worstDist()) {
            searchLevelExact<with_removed>(result_set, vec, otherChild, new_distsq, epsError);
        }
    }

    void addPointToTree(NodePtr node, int ind)
    {
        ElementType* point = points_[ind];

        if ((node->child1==NULL) && (node->child2==NULL)) {
            ElementType* leaf_point = node->point;
            ElementType max_span = 0;
            size_t div_feat = 0;
            for (size_t i=0;i<veclen_;++i) {
                ElementType span = std::abs(point[i]-leaf_point[i]);
                if (span > max_span) {
                    max_span = span;
                    div_feat = i;
                }
            }
            NodePtr left = new(pool_) Node();
            left->child1 = left->child2 = NULL;
            NodePtr right = new(pool_) Node();
            right->child1 = right->child2 = NULL;

            if (point[div_feat]<leaf_point[div_feat]) {
                left->divfeat = ind;
                left->point = point;
                right->divfeat = node->divfeat;
                right->point = node->point;
            }
            else {
                left->divfeat = node->divfeat;
                left->point = node->point;
                right->divfeat = ind;
                right->point = point;
            }
            node->divfeat = div_feat;
            node->divval = (point[div_feat]+leaf_point[div_feat])/2;
            node->child1 = left;
            node->child2 = right;
        }
        else {
            if (point[node->divfeat]<node->divval) {
                addPointToTree(node->child1,ind);
            }
            else {
                addPointToTree(node->child2,ind);
            }
        }
    }
private:
    void swap(KDTreeIndex& other)
    {
    	BaseClass::swap(other);
    	std::swap(trees_, other.trees_);
    	std::swap(tree_roots_, other.tree_roots_);
    	std::swap(pool_, other.pool_);
    }

private:

    enum
    {
        /**
         * To improve efficiency, only SAMPLE_MEAN random values are used to
         * compute the mean and variance at each level when building a tree.
         * A value of 100 seems to perform as well as using all values.
         */
        SAMPLE_MEAN = 100,
        /**
         * Top random dimensions to consider
         *
         * When creating random trees, the dimension on which to subdivide is
         * selected at random from among the top RAND_DIM dimensions with the
         * highest variance.  A value of 5 works well.
         */
        RAND_DIM=5
    };


    /**
     * Number of randomized trees that are used
     */
    int trees_;

    DistanceType* mean_;
    DistanceType* var_;

    /**
     * Array of k-d trees used to find neighbours.
     */
    std::vector<NodePtr> tree_roots_;

    /**
     * Pooled memory allocator.
     *
     * Using a pooled memory allocator is more efficient
     * than allocating memory directly when there is a large
     * number small of memory allocations.
     */
    PooledAllocator pool_;

    USING_BASECLASS_SYMBOLS
};   // class KDTreeIndex

}

#endif //FLANN_KDTREE_INDEX_H_

	/***********************************************************************
	* 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.
	*
	* 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.
	*************************************************************************/

	#ifndef FLANN_KDTREE_INDEX_H_
	#define FLANN_KDTREE_INDEX_H_

	#include <algorithm>
	#include <map>
	#include <cassert>
	#include <cstring>
	#include <stdarg.h>
	#include <cmath>
	#include <random>

	#include "FLANN/general.h"
	#include "FLANN/algorithms/nn_index.h"
	#include "FLANN/util/dynamic_bitset.h"
	#include "FLANN/util/matrix.h"
	#include "FLANN/util/result_set.h"
	#include "FLANN/util/heap.h"
	#include "FLANN/util/allocator.h"
	#include "FLANN/util/random.h"
	#include "FLANN/util/saving.h"


	namespace flann
	{

	struct KDTreeIndexParams : public IndexParams
	{
	KDTreeIndexParams(int trees = 4)
	{
	(*this)["algorithm"] = FLANN_INDEX_KDTREE;
	(*this)["trees"] = trees;
	}
	};


	/**
	* Randomized kd-tree index
	*
	* Contains the k-d trees and other information for indexing a set of points
	* for nearest-neighbor matching.
	*/
	template <typename Distance>
	class KDTreeIndex : public NNIndex<Distance>
	{
	public:
	typedef typename Distance::ElementType ElementType;
	typedef typename Distance::ResultType DistanceType;

	typedef NNIndex<Distance> BaseClass;

	typedef bool needs_kdtree_distance;


	/**
	* KDTree constructor
	*
	* Params:
	* inputData = dataset with the input features
	* params = parameters passed to the kdtree algorithm
	*/
	KDTreeIndex(const IndexParams& params = KDTreeIndexParams(), Distance d = Distance() ) :
	BaseClass(params, d), mean_(NULL), var_(NULL)
	{
	trees_ = get_param(index_params_,"trees",4);
	}


	/**
	* KDTree constructor
	*
	* Params:
	* inputData = dataset with the input features
	* params = parameters passed to the kdtree algorithm
	*/
	KDTreeIndex(const Matrix<ElementType>& dataset, const IndexParams& params = KDTreeIndexParams(),
	Distance d = Distance() ) : BaseClass(params,d ), mean_(NULL), var_(NULL)
	{
	trees_ = get_param(index_params_,"trees",4);

	setDataset(dataset);
	}

	KDTreeIndex(const KDTreeIndex& other) : BaseClass(other),
	trees_(other.trees_)
	{
	tree_roots_.resize(other.tree_roots_.size());
	for (size_t i=0;i<tree_roots_.size();++i) {
	copyTree(tree_roots_[i], other.tree_roots_[i]);
	}
	}

	KDTreeIndex& operator=(KDTreeIndex other)
	{
	this->swap(other);
	return *this;
	}

	/**
	* Standard destructor
	*/
	virtual ~KDTreeIndex()
	{
	freeIndex();
	}

	BaseClass* clone() const
	{
	return new KDTreeIndex(*this);
	}

	using BaseClass::buildIndex;

	void addPoints(const Matrix<ElementType>& points, float rebuild_threshold = 2)
	{
	assert(points.cols==veclen_);

	size_t old_size = size_;
	extendDataset(points);

	if (rebuild_threshold>1 && size_at_build_*rebuild_threshold<size_) {
	buildIndex();
	}
	else {
	for (size_t i=old_size;i<size_;++i) {
	for (int j = 0; j < trees_; j++) {
	addPointToTree(tree_roots_[j], i);
	}
	}
	}
	}

	flann_algorithm_t getType() const
	{
	return FLANN_INDEX_KDTREE;
	}


	template<typename Archive>
	void serialize(Archive& ar)
	{
	ar.setObject(this);

	ar & static_cast<NNIndex<Distance>>(this);

	ar & trees_;

	if (Archive::is_loading::value) {
	tree_roots_.resize(trees_);
	}
	for (size_t i=0;i<tree_roots_.size();++i) {
	if (Archive::is_loading::value) {
	tree_roots_[i] = new(pool_) Node();
	}
	ar & *tree_roots_[i];
	}

	if (Archive::is_loading::value) {
	index_params_["algorithm"] = getType();
	index_params_["trees"] = trees_;
	}
	}


	void saveIndex(FILE* stream)
	{
	serialization::SaveArchive sa(stream);
	sa & *this;
	}


	void loadIndex(FILE* stream)
	{
	freeIndex();
	serialization::LoadArchive la(stream);
	la & *this;
	}

	/**
	* Computes the inde memory usage
	* Returns: memory used by the index
	*/
	int usedMemory() const
	{
	return int(pool_.usedMemory+pool_.wastedMemory+size_*sizeof(int)); // pool memory and vind array memory
	}

	/**
	* Find set of nearest neighbors to vec. Their indices are stored inside
	* the result object.
	*
	* Params:
	* result = the result object in which the indices of the nearest-neighbors are stored
	* vec = the vector for which to search the nearest neighbors
	* maxCheck = the maximum number of restarts (in a best-bin-first manner)
	*/
	void findNeighbors(ResultSet<DistanceType>& result, const ElementType* vec, const SearchParams& searchParams) const
	{
	int maxChecks = searchParams.checks;
	float epsError = 1+searchParams.eps;

	if (maxChecks==FLANN_CHECKS_UNLIMITED) {
	if (removed_) {
	getExactNeighbors<true>(result, vec, epsError);
	}
	else {
	getExactNeighbors<false>(result, vec, epsError);
	}
	}
	else {
	if (removed_) {
	getNeighbors<true>(result, vec, maxChecks, epsError);
	}
	else {
	getNeighbors<false>(result, vec, maxChecks, epsError);
	}
	}
	}

	protected:

	/**
	* Builds the index
	*/
	void buildIndexImpl()
	{
	// Create a permutable array of indices to the input vectors.
	std::vector<int> ind(size_);
	for (size_t i = 0; i < size_; ++i) {
	ind[i] = int(i);
	}

	mean_ = new DistanceType[veclen_];
	var_ = new DistanceType[veclen_];

	std::default_random_engine generator;

	tree_roots_.resize(trees_);
	/* Construct the randomized trees. */
	for (int i = 0; i < trees_; i++) {
	/* Randomize the order of vectors to allow for unbiased sampling. */
	std::shuffle(ind.begin(), ind.end(), generator);
	tree_roots_[i] = divideTree(&ind[0], int(size_) );
	}
	delete[] mean_;
	delete[] var_;
	}

	void freeIndex()
	{
	for (size_t i=0;i<tree_roots_.size();++i) {
	// using placement new, so call destructor explicitly
	if (tree_roots_[i]!=NULL) tree_roots_[i]->~Node();
	}
	pool_.free();
	}


	private:

	/--------------------- Internal Data Structures --------------------------/
	struct Node
	{
	/**
	* Dimension used for subdivision.
	*/
	int divfeat;
	/**
	* The values used for subdivision.
	*/
	DistanceType divval;
	/**
	* Point data
	*/
	ElementType* point;
	/**
	* The child nodes.
	*/
	Node* child1, *child2;
	Node(){
	child1 = NULL;
	child2 = NULL;
	}
	~Node() {
	if (child1 != NULL) { child1->~Node(); child1 = NULL; }

	if (child2 != NULL) { child2->~Node(); child2 = NULL; }
	}

	private:
	template<typename Archive>
	void serialize(Archive& ar)
	{
	typedef KDTreeIndex<Distance> Index;
	Index* obj = static_cast<Index*>(ar.getObject());

	ar & divfeat;
	ar & divval;

	bool leaf_node = false;
	if (Archive::is_saving::value) {
	leaf_node = ((child1==NULL) && (child2==NULL));
	}
	ar & leaf_node;

	if (leaf_node) {
	if (Archive::is_loading::value) {
	point = obj->points_[divfeat];
	}
	}

	if (!leaf_node) {
	if (Archive::is_loading::value) {
	child1 = new(obj->pool_) Node();
	child2 = new(obj->pool_) Node();
	}
	ar & *child1;
	ar & *child2;
	}
	}
	friend struct serialization::access;
	};
	typedef Node* NodePtr;
	typedef BranchStruct<NodePtr, DistanceType> BranchSt;
	typedef BranchSt* Branch;


	void copyTree(NodePtr& dst, const NodePtr& src)
	{
	dst = new(pool_) Node();
	dst->divfeat = src->divfeat;
	dst->divval = src->divval;
	if (src->child1==NULL && src->child2==NULL) {
	dst->point = points_[dst->divfeat];
	dst->child1 = NULL;
	dst->child2 = NULL;
	}
	else {
	copyTree(dst->child1, src->child1);
	copyTree(dst->child2, src->child2);
	}
	}

	/**
	* Create a tree node that subdivides the list of vecs from vind[first]
	* to vind[last]. The routine is called recursively on each sublist.
	* Place a pointer to this new tree node in the location pTree.
	*
	* Params: pTree = the new node to create
	* first = index of the first vector
	* last = index of the last vector
	*/
	NodePtr divideTree(int* ind, int count)
	{
	NodePtr node = new(pool_) Node(); // allocate memory

	/* If too few exemplars remain, then make this a leaf node. */
	if (count == 1) {
	node->child1 = node->child2 = NULL; /* Mark as leaf node. */
	node->divfeat = ind; / Store index of this vec. */
	node->point = points_[*ind];
	}
	else {
	int idx;
	int cutfeat;
	DistanceType cutval;
	meanSplit(ind, count, idx, cutfeat, cutval);

	node->divfeat = cutfeat;
	node->divval = cutval;
	node->child1 = divideTree(ind, idx);
	node->child2 = divideTree(ind+idx, count-idx);
	}

	return node;
	}


	/**
	* Choose which feature to use in order to subdivide this set of vectors.
	* Make a random choice among those with the highest variance, and use
	* its variance as the threshold value.
	*/
	void meanSplit(int* ind, int count, int& index, int& cutfeat, DistanceType& cutval)
	{
	memset(mean_,0,veclen_*sizeof(DistanceType));
	memset(var_,0,veclen_*sizeof(DistanceType));

	/* Compute mean values. Only the first SAMPLE_MEAN values need to be
	sampled to get a good estimate.
	*/
	int cnt = std::min((int)SAMPLE_MEAN+1, count);
	for (int j = 0; j < cnt; ++j) {
	ElementType* v = points_[ind[j]];
	for (size_t k=0; k<veclen_; ++k) {
	mean_[k] += v[k];
	}
	}
	DistanceType div_factor = DistanceType(1)/cnt;
	for (size_t k=0; k<veclen_; ++k) {
	mean_[k] *= div_factor;
	}

	/* Compute variances (no need to divide by count). */
	for (int j = 0; j < cnt; ++j) {
	ElementType* v = points_[ind[j]];
	for (size_t k=0; k<veclen_; ++k) {
	DistanceType dist = v[k] - mean_[k];
	var_[k] += dist * dist;
	}
	}
	/* Select one of the highest variance indices at random. */
	cutfeat = selectDivision(var_);
	cutval = mean_[cutfeat];

	int lim1, lim2;
	planeSplit(ind, count, cutfeat, cutval, lim1, lim2);

	if (lim1>count/2) index = lim1;
	else if (lim2<count/2) index = lim2;
	else index = count/2;

	/* If either list is empty, it means that all remaining features
	* are identical. Split in the middle to maintain a balanced tree.
	*/
	if ((lim1==count)\|\|(lim2==0)) index = count/2;
	}


	/**
	* Select the top RAND_DIM largest values from v and return the index of
	* one of these selected at random.
	*/
	int selectDivision(DistanceType* v)
	{
	int num = 0;
	size_t topind[RAND_DIM];

	/* Create a list of the indices of the top RAND_DIM values. */
	for (size_t i = 0; i < veclen_; ++i) {
	if ((num < RAND_DIM)\|\|(v[i] > v[topind[num-1]])) {
	/* Put this element at end of topind. */
	if (num < RAND_DIM) {
	topind[num++] = i; /* Add to list. */
	}
	else {
	topind[num-1] = i; /* Replace last element. */
	}
	/* Bubble end value down to right location by repeated swapping. */
	int j = num - 1;
	while (j > 0 && v[topind[j]] > v[topind[j-1]]) {
	std::swap(topind[j], topind[j-1]);
	--j;
	}
	}
	}
	/* Select a random integer in range [0,num-1], and return that index. */
	int rnd = rand_int(num);
	return (int)topind[rnd];
	}


	/**
	* Subdivide the list of points by a plane perpendicular on axe corresponding
	* to the 'cutfeat' dimension at 'cutval' position.
	*
	* On return:
	* dataset[ind[0..lim1-1]][cutfeat]<cutval
	* dataset[ind[lim1..lim2-1]][cutfeat]==cutval
	* dataset[ind[lim2..count]][cutfeat]>cutval
	*/
	void planeSplit(int* ind, int count, int cutfeat, DistanceType cutval, int& lim1, int& lim2)
	{
	/* Move vector indices for left subtree to front of list. */
	int left = 0;
	int right = count-1;
	for (;; ) {
	while (left<=right && points_[ind[left]][cutfeat]<cutval) ++left;
	while (left<=right && points_[ind[right]][cutfeat]>=cutval) --right;
	if (left>right) break;
	std::swap(ind[left], ind[right]); ++left; --right;
	}
	lim1 = left;
	right = count-1;
	for (;; ) {
	while (left<=right && points_[ind[left]][cutfeat]<=cutval) ++left;
	while (left<=right && points_[ind[right]][cutfeat]>cutval) --right;
	if (left>right) break;
	std::swap(ind[left], ind[right]); ++left; --right;
	}
	lim2 = left;
	}

	/**
	* Performs an exact nearest neighbor search. The exact search performs a full
	* traversal of the tree.
	*/
	template<bool with_removed>
	void getExactNeighbors(ResultSet<DistanceType>& result, const ElementType* vec, float epsError) const
	{
	// checkID -= 1; /* Set a different unique ID for each search. */

	if (trees_ > 1) {
	fprintf(stderr,"It doesn't make any sense to use more than one tree for exact search");
	}
	if (trees_>0) {
	searchLevelExact<with_removed>(result, vec, tree_roots_[0], 0.0, epsError);
	}
	}

	/**
	* Performs the approximate nearest-neighbor search. The search is approximate
	* because the tree traversal is abandoned after a given number of descends in
	* the tree.
	*/
	template<bool with_removed>
	void getNeighbors(ResultSet<DistanceType>& result, const ElementType* vec, int maxCheck, float epsError) const
	{
	int i;
	BranchSt branch;

	int checkCount = 0;
	Heap<BranchSt>* heap = new Heap<BranchSt>((int)size_);
	DynamicBitset checked(size_);

	/* Search once through each tree down to root. */
	for (i = 0; i < trees_; ++i) {
	searchLevel<with_removed>(result, vec, tree_roots_[i], 0, checkCount, maxCheck, epsError, heap, checked);
	}

	/* Keep searching other branches from heap until finished. */
	while ( heap->popMin(branch) && (checkCount < maxCheck \|\| !result.full() )) {
	searchLevel<with_removed>(result, vec, branch.node, branch.mindist, checkCount, maxCheck, epsError, heap, checked);
	}

	delete heap;

	}

	/**
	* Search starting from a given node of the tree. Based on any mismatches at
	* higher levels, all exemplars below this level must have a distance of
	* at least "mindistsq".
	*/
	template<bool with_removed>
	void searchLevel(ResultSet<DistanceType>& result_set, const ElementType* vec, NodePtr node, DistanceType mindist, int& checkCount, int maxCheck,
	float epsError, Heap<BranchSt>* heap, DynamicBitset& checked) const
	{
	if (result_set.worstDist()<mindist) {
	// printf("Ignoring branch, too far\n");
	return;
	}

	/* If this is a leaf node, then do check and return. */
	if ((node->child1 == NULL)&&(node->child2 == NULL)) {
	int index = node->divfeat;
	if (with_removed) {
	if (removed_points_.test(index)) return;
	}
	/* Do not check same node more than once when searching multiple trees. */
	if ( checked.test(index) \|\| ((checkCount>=maxCheck)&& result_set.full()) ) return;
	checked.set(index);
	checkCount++;

	DistanceType dist = distance_(node->point, vec, veclen_);
	result_set.addPoint(dist,index);
	return;
	}

	/* Which child branch should be taken first? */
	ElementType val = vec[node->divfeat];
	DistanceType diff = val - node->divval;
	NodePtr bestChild = (diff < 0) ? node->child1 : node->child2;
	NodePtr otherChild = (diff < 0) ? node->child2 : node->child1;

	/* Create a branch record for the branch not taken. Add distance
	of this feature boundary (we don't attempt to correct for any
	use of this feature in a parent node, which is unlikely to
	happen and would have only a small effect). Don't bother
	adding more branches to heap after halfway point, as cost of
	adding exceeds their value.
	*/

	DistanceType new_distsq = mindist + distance_.accum_dist(val, node->divval, node->divfeat);
	// if (2 * checkCount < maxCheck \|\| !result.full()) {
	if ((new_distsq*epsError < result_set.worstDist())\|\| !result_set.full()) {
	heap->insert( BranchSt(otherChild, new_distsq) );
	}

	/* Call recursively to search next level down. */
	searchLevel<with_removed>(result_set, vec, bestChild, mindist, checkCount, maxCheck, epsError, heap, checked);
	}

	/**
	* Performs an exact search in the tree starting from a node.
	*/
	template<bool with_removed>
	void searchLevelExact(ResultSet<DistanceType>& result_set, const ElementType* vec, const NodePtr node, DistanceType mindist, const float epsError) const
	{
	/* If this is a leaf node, then do check and return. */
	if ((node->child1 == NULL)&&(node->child2 == NULL)) {
	int index = node->divfeat;
	if (with_removed) {
	if (removed_points_.test(index)) return; // ignore removed points
	}
	DistanceType dist = distance_(node->point, vec, veclen_);
	result_set.addPoint(dist,index);

	return;
	}

	/* Which child branch should be taken first? */
	ElementType val = vec[node->divfeat];
	DistanceType diff = val - node->divval;
	NodePtr bestChild = (diff < 0) ? node->child1 : node->child2;
	NodePtr otherChild = (diff < 0) ? node->child2 : node->child1;

	/* Create a branch record for the branch not taken. Add distance
	of this feature boundary (we don't attempt to correct for any
	use of this feature in a parent node, which is unlikely to
	happen and would have only a small effect). Don't bother
	adding more branches to heap after halfway point, as cost of
	adding exceeds their value.
	*/

	DistanceType new_distsq = mindist + distance_.accum_dist(val, node->divval, node->divfeat);

	/* Call recursively to search next level down. */
	searchLevelExact<with_removed>(result_set, vec, bestChild, mindist, epsError);

	if (mindist*epsError<=result_set.worstDist()) {
	searchLevelExact<with_removed>(result_set, vec, otherChild, new_distsq, epsError);
	}
	}

	void addPointToTree(NodePtr node, int ind)
	{
	ElementType* point = points_[ind];

	if ((node->child1==NULL) && (node->child2==NULL)) {
	ElementType* leaf_point = node->point;
	ElementType max_span = 0;
	size_t div_feat = 0;
	for (size_t i=0;i<veclen_;++i) {
	ElementType span = std::abs(point[i]-leaf_point[i]);
	if (span > max_span) {
	max_span = span;
	div_feat = i;
	}
	}
	NodePtr left = new(pool_) Node();
	left->child1 = left->child2 = NULL;
	NodePtr right = new(pool_) Node();
	right->child1 = right->child2 = NULL;

	if (point[div_feat]<leaf_point[div_feat]) {
	left->divfeat = ind;
	left->point = point;
	right->divfeat = node->divfeat;
	right->point = node->point;
	}
	else {
	left->divfeat = node->divfeat;
	left->point = node->point;
	right->divfeat = ind;
	right->point = point;
	}
	node->divfeat = div_feat;
	node->divval = (point[div_feat]+leaf_point[div_feat])/2;
	node->child1 = left;
	node->child2 = right;
	}
	else {
	if (point[node->divfeat]<node->divval) {
	addPointToTree(node->child1,ind);
	}
	else {
	addPointToTree(node->child2,ind);
	}
	}
	}
	private:
	void swap(KDTreeIndex& other)
	{
	BaseClass::swap(other);
	std::swap(trees_, other.trees_);
	std::swap(tree_roots_, other.tree_roots_);
	std::swap(pool_, other.pool_);
	}

	private:

	enum
	{
	/**
	* To improve efficiency, only SAMPLE_MEAN random values are used to
	* compute the mean and variance at each level when building a tree.
	* A value of 100 seems to perform as well as using all values.
	*/
	SAMPLE_MEAN = 100,
	/**
	* Top random dimensions to consider
	*
	* When creating random trees, the dimension on which to subdivide is
	* selected at random from among the top RAND_DIM dimensions with the
	* highest variance. A value of 5 works well.
	*/
	RAND_DIM=5
	};


	/**
	* Number of randomized trees that are used
	*/
	int trees_;

	DistanceType* mean_;
	DistanceType* var_;

	/**
	* Array of k-d trees used to find neighbours.
	*/
	std::vector<NodePtr> tree_roots_;

	/**
	* Pooled memory allocator.
	*
	* Using a pooled memory allocator is more efficient
	* than allocating memory directly when there is a large
	* number small of memory allocations.
	*/
	PooledAllocator pool_;

	USING_BASECLASS_SYMBOLS
	}; // class KDTreeIndex

	}

	#endif //FLANN_KDTREE_INDEX_H_