#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);

}

{

    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];

    int kball = kdtree.annkFRSearch(annq, k, sqRad, nn_idx, dists, eps);

    bp::list indices, distances;

    for (unsigned i = 0; i < k; ++i)

    {

        if (nn_idx[i] != ANN_NULL_IDX)

            indices.append(nn_idx[i]);

        if (dists[i] != ANN_DIST_INF)

            distances.append(dists[i]);

    }

    delete nn_idx; delete dists;

    return bp::make_tuple(indices, distances, kball);

}

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("kPriSearch",          &k_priority_search)

        .def("kFRSearch",           &k_fixed_radius_search)

        .def("__len__",             &ANNkd_tree::nPoints)

        .def("dim",                 &ANNkd_tree::theDim)

    ;

    bp::def("max_pts_visit",        &annMaxPtsVisit);

}