<|start_filename|>src/stream_compaction.h<|end_filename|>
#pragma once

#include <tbb/parallel_for.h>
#include <tbb/parallel_scan.h>

#include <functional>

namespace pointcloud_cropper {

// General purpose parallelized stream compaction with TBB.
template <typename dtype_t, typename index_t, typename func_t>
index_t GetValidIndices(const dtype_t* src,
                        index_t* valid_indices,
                        index_t N,
                        func_t is_valid) {
    index_t num_valid = tbb::parallel_scan(
            tbb::blocked_range<index_t>(0, N), 0,
            [&](const tbb::blocked_range<index_t>& r, index_t sum,
                bool is_final_scan) -> index_t {
                index_t temp = sum;
                for (index_t i = r.begin(); i < r.end(); ++i) {
                    if (is_valid(src[i])) {
                        if (is_final_scan) {
                            valid_indices[temp] = i;
                        }
                        temp += 1;
                    }
                }
                return temp;
            },
            [](int left, int right) { return left + right; });
    return num_valid;
}

// General purpose stream compaction serial implementation.
template <typename dtype_t, typename index_t, typename func_t>
index_t GetValidIndicesSerial(const dtype_t* src,
                              index_t* valid_indices,
                              index_t N,
                              func_t is_valid) {
    index_t num_valid = 0;
    for (index_t i = 0; i < N; i++) {
        if (is_valid(src[i])) {
            valid_indices[num_valid++] = i;
        }
    }
    return num_valid;
}

}  // namespace pointcloud_cropper


<|start_filename|>tbb.cmake<|end_filename|>
include(ExternalProject)

ExternalProject_Add(
    ext_tbb
    PREFIX tbb
    GIT_REPOSITORY https://github.com/wjakob/tbb.git
    GIT_TAG 141b0e310e1fb552bdca887542c9c1a8544d6503 # Sept 2020
    UPDATE_COMMAND ""
    CMAKE_ARGS
        -DCMAKE_INSTALL_PREFIX=<INSTALL_DIR>
        -DCMAKE_POSITION_INDEPENDENT_CODE=ON
        -DTBB_BUILD_TBBMALLOC=ON
        -DTBB_BUILD_TBBMALLOC_PROXYC=OFF
        -DTBB_BUILD_SHARED=OFF
        -DTBB_BUILD_STATIC=ON
        -DTBB_BUILD_TESTS=OFF
        -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
)

ExternalProject_Get_property(ext_tbb INSTALL_DIR)
set(TBB_INCLUDE_DIR "${INSTALL_DIR}/include")
set(TBB_LIB_DIR "${INSTALL_DIR}/lib")
set(TBB_LIBRARIES tbb_static tbbmalloc_static)


<|start_filename|>src/pointcloud_cropper.cpp<|end_filename|>
#include <pybind11/functional.h>
#include <pybind11/numpy.h>
#include <pybind11/operators.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/stl_bind.h>

#include "stream_compaction.h"

namespace pointcloud_cropper {

namespace py = pybind11;
using namespace py::literals;

template <class T>
struct Point {
    T vals[3];
};

static py::array ToIntArray1D(const int *array_buf, int N) {
    py::dtype py_dtype(py::format_descriptor<int32_t>::format());
    py::array::ShapeContainer py_shape({N});
    py::array::StridesContainer py_strides({4});  // int32_t is 4 bytes

    auto py_destructor = [](PyObject *data) {
        // Deallocates buffer when the numpy variable goes out of scope.
        int *buf = reinterpret_cast<int *>(
                PyCapsule_GetPointer(data, "int32_t buffer"));
        if (buf) {
            // std::cout << "destructed" << std::endl;
            delete[] buf;
        } else {
            PyErr_Clear();
        }
    };
    py::capsule py_capsule(array_buf, "int32_t buffer", py_destructor);
    return py::array(py_dtype, py_shape, py_strides, array_buf, py_capsule);
}

template <typename T>
static py::array GetCropIndex(py::array_t<T> points,
                              double x_min,
                              double y_min,
                              double z_min,
                              double x_max,
                              double y_max,
                              double z_max) {
    // Parse input buffer as array of Point<T>.
    py::buffer_info info = points.request();
    std::vector<int> shape(info.shape.begin(), info.shape.end());
    if (shape.size() != 2 || shape[1] != 3) {
        throw std::runtime_error("points must be (N, 3)");
    }
    int N = shape[0];
    const Point<T> *points_ptr = reinterpret_cast<const Point<T> *>(info.ptr);

    // Output buffer will be freed by py::capsule.
    int *valid_indices = new int[N];

    // Perform stream compaction.
    auto is_valid = [&](const Point<T> &point) -> bool {
        return point.vals[0] >= x_min && point.vals[0] <= x_max &&
               point.vals[1] >= y_min && point.vals[1] <= y_max &&
               point.vals[2] >= z_min && point.vals[2] <= z_max;
    };
    int num_valid = GetValidIndices(points_ptr, valid_indices, N, is_valid);
    return ToIntArray1D(valid_indices, num_valid);
}

template <typename T>
static py::array GetCropIndexSerial(py::array_t<T> points,
                                    double x_min,
                                    double y_min,
                                    double z_min,
                                    double x_max,
                                    double y_max,
                                    double z_max) {
    // Parse input buffer as array of Point<T>.
    py::buffer_info info = points.request();
    std::vector<int> shape(info.shape.begin(), info.shape.end());
    if (shape.size() != 2 || shape[1] != 3) {
        throw std::runtime_error("points must be (N, 3)");
    }
    int N = shape[0];
    const Point<T> *points_ptr = reinterpret_cast<const Point<T> *>(info.ptr);

    // Output buffer will be freed by py::capsule.
    int *valid_indices = new int[N];

    // Perform stream compaction.
    auto is_valid = [&](const Point<T> &point) -> bool {
        return point.vals[0] >= x_min && point.vals[0] <= x_max &&
               point.vals[1] >= y_min && point.vals[1] <= y_max &&
               point.vals[2] >= z_min && point.vals[2] <= z_max;
    };
    int num_valid =
            GetValidIndicesSerial(points_ptr, valid_indices, N, is_valid);
    return ToIntArray1D(valid_indices, num_valid);
}

PYBIND11_MODULE(pointcloud_cropper, m) {
    m.def("get_crop_indices", &GetCropIndex<double>, "points"_a, "x_min"_a,
          "y_min"_a, "z_min"_a, "x_max"_a, "y_max"_a, "z_max"_a);
    m.def("get_crop_indices", &GetCropIndex<float>, "points"_a, "x_min"_a,
          "y_min"_a, "z_min"_a, "x_max"_a, "y_max"_a, "z_max"_a);
    m.def("get_crop_indices_serial", &GetCropIndexSerial<double>, "points"_a,
          "x_min"_a, "y_min"_a, "z_min"_a, "x_max"_a, "y_max"_a, "z_max"_a);
    m.def("get_crop_indices_serial", &GetCropIndexSerial<float>, "points"_a,
          "x_min"_a, "y_min"_a, "z_min"_a, "x_max"_a, "y_max"_a, "z_max"_a);
}

}  // namespace pointcloud_cropper
