#ifndef AKNN_H
#define AKNN_H

#include "nanoflann.hpp"

using namespace nanoflann;

template <typename T, int dim> class aKNN {

	struct DataSet {

		std::vector<T> elems;

		inline size_t kdtree_get_point_count() const {return elems.size();}
		inline float kdtree_distance(const float* p1, const size_t idxP2, size_t) const {
			float dist = 0;
			for (int i = 0; i < dim; ++i) {
				float delta = (p1[i] - kdtree_get_pt(idxP2, i));
				dist += delta*delta;
			}
			return dist;
		}

		inline float kdtree_get_pt(const size_t idx, int pos) const {
			return elems[idx].feature[pos];
		}

		template <class BBOX> bool kdtree_get_bbox(BBOX&) const {return false;}

	} data;

	typedef KDTreeSingleIndexAdaptor<L2_Simple_Adaptor<float, DataSet>, DataSet, dim> MyTree;

	MyTree* tree = nullptr;



public:

	/** add a new element */
	void add(const T& elem) {
		data.elems.push_back(elem);
	}

	/** build the KD-Tree */
	void build() {
		tree = new MyTree(dim, data, KDTreeSingleIndexAdaptorParams(10) );
		tree->buildIndex();
	}

	/** get the nearest n elements */
	template <typename T2> std::vector<T> get(const T2* query, const int numResults) {

		float distances[numResults];
		size_t indices[numResults];

		KNNResultSet<float> res(numResults);
		res.init(indices, distances);

		tree->knnSearch(query, numResults, indices, distances);
		std::vector<T> vec;
		for (int i = 0; i < numResults; ++i) {
			vec.push_back(data.elems[indices[i]]);
		}
		return vec;

	}


};

#endif // AKNN_H