// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2009 Ilya Baran <ibaran@mit.edu>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0

#ifndef KDBVH_H_INCLUDED
#define KDBVH_H_INCLUDED

// IWYU pragma: private
#include "./InternalHeaderCheck.h"

namespace Eigen {

namespace internal {

// internal pair class for the BVH--used instead of std::pair because of alignment
template <typename Scalar, int Dim>
struct vector_int_pair {
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW_IF_VECTORIZABLE_FIXED_SIZE(Scalar, Dim)
  typedef Matrix<Scalar, Dim, 1> VectorType;

  vector_int_pair(const VectorType &v, int i) : first(v), second(i) {}

  VectorType first;
  int second;
};

// these templates help the tree initializer get the bounding boxes either from a provided
// iterator range or using bounding_box in a unified way
template <typename ObjectList, typename VolumeList, typename BoxIter>
struct get_boxes_helper {
  void operator()(const ObjectList &objects, BoxIter boxBegin, BoxIter boxEnd, VolumeList &outBoxes) const {
    outBoxes.insert(outBoxes.end(), boxBegin, boxEnd);
    eigen_assert(outBoxes.size() == objects.size());
    EIGEN_ONLY_USED_FOR_DEBUG(objects);
  }
};

template <typename ObjectList, typename VolumeList>
struct get_boxes_helper<ObjectList, VolumeList, int> {
  void operator()(const ObjectList &objects, int, int, VolumeList &outBoxes) const {
    outBoxes.reserve(objects.size());
    for (int i = 0; i < (int)objects.size(); ++i) outBoxes.push_back(bounding_box(objects[i]));
  }
};

}  // end namespace internal

/** \class KdBVH
 *  \brief A simple bounding volume hierarchy based on AlignedBox
 *
 *  \param Scalar_ The underlying scalar type of the bounding boxes
 *  \param Dim_ The dimension of the space in which the hierarchy lives
 *  \param Object_ The object type that lives in the hierarchy.  It must have value semantics.  Either
 *                 `bounding_box(Object_)` must be defined and return an `AlignedBox<Scalar_, Dim_>` or bounding boxes
 *                  must be provided to the tree initializer.
 *
 * This class provides a simple (as opposed to optimized) implementation of a bounding volume hierarchy analogous to a
 * Kd-tree. Given a sequence of objects, it computes their bounding boxes, constructs a Kd-tree of their centers and
 * builds a BVH with the structure of that Kd-tree.  When the elements of the tree are too expensive to be copied
 * around, it is useful for `Object_` to be a pointer.
 */
template <typename Scalar_, int Dim_, typename Object_>
class KdBVH {
 public:
  enum { Dim = Dim_ };
  typedef Object_ Object;
  typedef std::vector<Object, aligned_allocator<Object> > ObjectList;
  typedef Scalar_ Scalar;
  typedef AlignedBox<Scalar, Dim> Volume;
  typedef std::vector<Volume, aligned_allocator<Volume> > VolumeList;
  typedef int Index;
  typedef const int *VolumeIterator;  // the iterators are just pointers into the tree's vectors
  typedef const Object *ObjectIterator;

  KdBVH() = default;

  /** Given an iterator range over \a Object references, constructs the BVH.  Requires that bounding_box(Object) return
   * a Volume. */
  template <typename Iter>
  KdBVH(Iter begin, Iter end) {
    init(begin, end, 0, 0);
  }  // int is recognized by init as not being an iterator type

  /** Given an iterator range over \a Object references and an iterator range over their bounding boxes, constructs the
   * BVH */
  template <typename OIter, typename BIter>
  KdBVH(OIter begin, OIter end, BIter boxBegin, BIter boxEnd) {
    init(begin, end, boxBegin, boxEnd);
  }

  /** Given an iterator range over \a Object references, constructs the BVH, overwriting whatever is in there currently.
   * Requires that bounding_box(Object) return a Volume. */
  template <typename Iter>
  void init(Iter begin, Iter end) {
    init(begin, end, 0, 0);
  }

  /** Given an iterator range over \a Object references and an iterator range over their bounding boxes,
   * constructs the BVH, overwriting whatever is in there currently. */
  template <typename OIter, typename BIter>
  void init(OIter begin, OIter end, BIter boxBegin, BIter boxEnd) {
    objects.clear();
    boxes.clear();
    children.clear();

    objects.insert(objects.end(), begin, end);
    int n = static_cast<int>(objects.size());

    if (n < 2) return;  // if we have at most one object, we don't need any internal nodes

    VolumeList objBoxes;
    VIPairList objCenters;

    // compute the bounding boxes depending on BIter type
    internal::get_boxes_helper<ObjectList, VolumeList, BIter>()(objects, boxBegin, boxEnd, objBoxes);

    objCenters.reserve(n);
    boxes.reserve(n - 1);
    children.reserve(2 * n - 2);

    for (int i = 0; i < n; ++i) objCenters.push_back(VIPair(objBoxes[i].center(), i));

    build(objCenters, 0, n, objBoxes, 0);  // the recursive part of the algorithm

    ObjectList tmp(n);
    tmp.swap(objects);
    for (int i = 0; i < n; ++i) objects[i] = tmp[objCenters[i].second];
  }

  /** \returns the index of the root of the hierarchy */
  inline Index getRootIndex() const { return (int)boxes.size() - 1; }

  /** Given an \a index of a node, on exit, \a outVBegin and \a outVEnd range over the indices of the volume children of
   * the node and \a outOBegin and \a outOEnd range over the object children of the node */
  EIGEN_STRONG_INLINE void getChildren(Index index, VolumeIterator &outVBegin, VolumeIterator &outVEnd,
                                       ObjectIterator &outOBegin, ObjectIterator &outOEnd)
      const {  // inlining this function should open lots of optimization opportunities to the compiler
    if (index < 0) {
      outVBegin = outVEnd;
      if (!objects.empty()) outOBegin = &(objects[0]);
      outOEnd = outOBegin + objects.size();  // output all objects--necessary when the tree has only one object
      return;
    }

    int numBoxes = static_cast<int>(boxes.size());

    int idx = index * 2;
    if (children[idx + 1] < numBoxes) {  // second index is always bigger
      outVBegin = &(children[idx]);
      outVEnd = outVBegin + 2;
      outOBegin = outOEnd;
    } else if (children[idx] >= numBoxes) {  // if both children are objects
      outVBegin = outVEnd;
      outOBegin = &(objects[children[idx] - numBoxes]);
      outOEnd = outOBegin + 2;
    } else {  // if the first child is a volume and the second is an object
      outVBegin = &(children[idx]);
      outVEnd = outVBegin + 1;
      outOBegin = &(objects[children[idx + 1] - numBoxes]);
      outOEnd = outOBegin + 1;
    }
  }

  /** \returns the bounding box of the node at \a index */
  inline const Volume &getVolume(Index index) const { return boxes[index]; }

 private:
  typedef internal::vector_int_pair<Scalar, Dim> VIPair;
  typedef std::vector<VIPair, aligned_allocator<VIPair> > VIPairList;
  typedef Matrix<Scalar, Dim, 1> VectorType;
  struct VectorComparator  // compares vectors, or more specifically, VIPairs along a particular dimension
  {
    VectorComparator(int inDim) : dim(inDim) {}
    inline bool operator()(const VIPair &v1, const VIPair &v2) const { return v1.first[dim] < v2.first[dim]; }
    int dim;
  };

  // Build the part of the tree between objects[from] and objects[to] (not including objects[to]).
  // This routine partitions the objCenters in [from, to) along the dimension dim, recursively constructs
  // the two halves, and adds their parent node.  TODO: a cache-friendlier layout
  void build(VIPairList &objCenters, int from, int to, const VolumeList &objBoxes, int dim) {
    eigen_assert(to - from > 1);
    if (to - from == 2) {
      boxes.push_back(objBoxes[objCenters[from].second].merged(objBoxes[objCenters[from + 1].second]));
      children.push_back(from + (int)objects.size() - 1);  // there are objects.size() - 1 tree nodes
      children.push_back(from + (int)objects.size());
    } else if (to - from == 3) {
      int mid = from + 2;
      std::nth_element(objCenters.begin() + from, objCenters.begin() + mid, objCenters.begin() + to,
                       VectorComparator(dim));  // partition
      build(objCenters, from, mid, objBoxes, (dim + 1) % Dim);
      int idx1 = (int)boxes.size() - 1;
      boxes.push_back(boxes[idx1].merged(objBoxes[objCenters[mid].second]));
      children.push_back(idx1);
      children.push_back(mid + (int)objects.size() - 1);
    } else {
      int mid = from + (to - from) / 2;
      nth_element(objCenters.begin() + from, objCenters.begin() + mid, objCenters.begin() + to,
                  VectorComparator(dim));  // partition
      build(objCenters, from, mid, objBoxes, (dim + 1) % Dim);
      int idx1 = (int)boxes.size() - 1;
      build(objCenters, mid, to, objBoxes, (dim + 1) % Dim);
      int idx2 = (int)boxes.size() - 1;
      boxes.push_back(boxes[idx1].merged(boxes[idx2]));
      children.push_back(idx1);
      children.push_back(idx2);
    }
  }

  std::vector<int> children;  // children of x are children[2x] and children[2x+1], indices bigger than boxes.size()
                              // index into objects.
  VolumeList boxes;
  ObjectList objects;
};

}  // end namespace Eigen

#endif  // KDBVH_H_INCLUDED