#include <boost/python.hpp>
#include <boost/python/stl_iterator.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/functional/hash.hpp>
namespace bp = boost::python;
#include <ANN/ANN.h>
#include <iostream>
// Constructor from list TODO: change to iterator
boost::shared_ptr<ANNkd_tree> init_from_list(bp::list lst)
{
int dimension = bp::len(lst[0]);
int npts = bp::len(lst);
ANNpointArray dataPts = annAllocPts(npts, dimension);
// Convert points from Python list to ANNpointArray
for (unsigned p = 0; p < bp::len(lst); ++p)
{
ANNpoint& pt = dataPts[p];
for (unsigned c = 0; c < dimension; ++c)
pt[c] = bp::extract<ANNcoord>(lst[p][c]);
}
boost::shared_ptr<ANNkd_tree> p(new ANNkd_tree(dataPts, npts, dimension));
return p;
}
bp::tuple search(ANNkd_tree& kdtree, bp::list q, int k, double eps, bool priority = false)
{
ANNpoint annq = annAllocPt(bp::len(q));
for (unsigned c = 0; c < bp::len(q); ++c)
annq[c] = bp::extract<ANNcoord>(q[c]);
ANNidxArray nn_idx = new ANNidx[k];
ANNdistArray dists = new ANNdist[k];
if (priority)
kdtree.annkPriSearch(annq, k, nn_idx, dists, eps);
else
kdtree.annkSearch(annq, k, nn_idx, dists, eps);
bp::list indices, distances;
for (unsigned i = 0; i < k; ++i)
{
indices.append(nn_idx[i]);
distances.append(dists[i]);
}
delete nn_idx; delete dists;
return bp::make_tuple(indices, distances);
}
bp::tuple ksearch(ANNkd_tree& kdtree, bp::list q, int k, double eps)
{
return search(kdtree, q, k, eps, false);
}
bp::tuple k_priority_search(ANNkd_tree& kdtree, bp::list q, int k, double eps)
{
return search(kdtree, q, k, eps, true);
}
void export_ann_kd_tree()
{
bp::class_<ANNkd_tree>("KDTree")
.def("__init__", bp::make_constructor(&init_from_list))
.def("ksearch", &ksearch)
.def("k_priority_search", &k_priority_search)
.def("__len__", &ANNkd_tree::nPoints)
.def("dim", &ANNkd_tree::theDim)
;
bp::def("max_pts_visit", &annMaxPtsVisit);
}