#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                           k_fixed_radius_search(ANNkd_tree& kdtree, bp::list q, double sqRad, int k, double eps)
{
    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);
}