#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                           ksearch(ANNkd_tree& kdtree, bp::list q, int k, double eps)
{
    ANNpointArray pts = kdtree.thePoints();
    // for (unsigned i = 0; i < kdtree.nPoints(); ++i)
    //     std::cout << pts[i][0] << " " << pts[i][1] << " " << pts[i][2] << std::endl;

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

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


void export_ann_kd_tree()
{
    bp::class_<ANNkd_tree>("KDTree")
        .def("__init__",            bp::make_constructor(&init_from_list))

        .def("ksearch",             &ksearch)
        .def("__len__",             &ANNkd_tree::nPoints)
        .def("dim",                 &ANNkd_tree::theDim)
    ;
}