text: // ----------------------------------------------------------------------------
  // -                        Open3D: www.open3d.org                            -
  // ----------------------------------------------------------------------------
  // The MIT License (MIT)
  //
  // Copyright (c) 2020 www.open3d.org
  //
  // Permission is hereby granted, free of charge, to any person obtaining a copy
  // of this software and associated documentation files (the "Software"), to deal
  // in the Software without restriction, including without limitation the rights
  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  // copies of the Software, and to permit persons to whom the Software is
  // furnished to do so, subject to the following conditions:
  //
  // The above copyright notice and this permission notice shall be included in
  // all copies or substantial portions of the Software.
  //
  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
  // IN THE SOFTWARE.
  // ----------------------------------------------------------------------------
  
  #pragma once
  
  #include <unordered_map>
  #include "Eigen/Core"
  #include "Open3D/Utility/Helper.h"
  #include "tbb/task_group.h"
  
  namespace open3d {
  namespace ml {
  namespace detail {
  
  namespace {
  
  enum AccumulationFn { AVERAGE = 0, NEAREST_NEIGHBOR, MAX, CENTER };
  
  template <class TReal,
            class TFeat,
            AccumulationFn POS_FN,
            AccumulationFn FEAT_FN>
  class Accumulator {
  public:
      Accumulator()
          : count_(0),
            min_sqr_dist_to_center_(std::numeric_limits<TReal>::max()),
            position_(0, 0, 0) {
          static_assert(POS_FN != MAX, "MAX is not allowed for point positions");
          static_assert(FEAT_FN != CENTER,
                        "CENTER is not allowed for feature vectors");
      }
  
      template <class Derived, class Derived2, class Derived3>
      inline void AddPoint(const Eigen::MatrixBase<Derived>& pos,
                           const Eigen::MatrixBase<Derived2>& voxel_center,
                           const Eigen::ArrayBase<Derived3>& feat) {
          bool new_nearest_neighbor = false;
          TReal sqr_d = 0;
          if (POS_FN == NEAREST_NEIGHBOR || FEAT_FN == NEAREST_NEIGHBOR) {
              sqr_d = (voxel_center - pos).squaredNorm();
              if (sqr_d < min_sqr_dist_to_center_) {
                  new_nearest_neighbor = true;
                  min_sqr_dist_to_center_ = sqr_d;
              }
          }
          if (POS_FN == AVERAGE) {
              position_ += pos.array();
          } else if (POS_FN == NEAREST_NEIGHBOR && new_nearest_neighbor) {
              position_ = pos;
          } else if (POS_FN == CENTER) {
              if (count_ == 0) position_ = voxel_center;
          }
  
          if (count_ == 0) {
              features_.resizeLike(feat);
              features_.setZero();
          }
          if (FEAT_FN == AVERAGE) {
              features_ += feat;
          } else if (FEAT_FN == NEAREST_NEIGHBOR && new_nearest_neighbor) {
              features_ = feat;
          } else if (FEAT_FN == MAX) {
              features_ = features_.max(feat);
          }
          ++count_;
      }
  
      inline Eigen::Array<TReal, 3, 1> Position() const {
          if (POS_FN == AVERAGE) {
              return position_ / count_;
          } else  // if( POS_FN == NEAREST_NEIGHBOR || POS_FN == CENTER )
          {
              return position_;
          }
      }
  
      inline Eigen::Array<TFeat, Eigen::Dynamic, 1> Features() const {
          if (FEAT_FN == AVERAGE) {
              return features_ / count_;
          } else  // if( FEAT_FN == NEAREST_NEIGHBOR || FEAT_FN == MAX )
          {
              return features_;
          }
      }
  
      inline int Count() const { return count_; }
  
  private:
      int count_;
      TReal min_sqr_dist_to_center_;
      Eigen::Array<TReal, 3, 1> position_;
      Eigen::Array<TFeat, Eigen::Dynamic, 1> features_;
  };  // namespace
  
  template <class TReal,
            class TFeat,
            AccumulationFn POS_FN,
            AccumulationFn FEAT_FN>
  class AccumulatorBackprop {
  public:
      AccumulatorBackprop()
          : count_(0),
            min_sqr_dist_to_center_(std::numeric_limits<TReal>::max()),
            position_(0, 0, 0) {
          static_assert(POS_FN != MAX, "MAX is not allowed for point positions");
          static_assert(FEAT_FN != CENTER,
                        "CENTER is not allowed for feature vectors");
      }
  
      template <class Derived, class Derived2, class Derived3>
      inline void AddPoint(const Eigen::MatrixBase<Derived>& pos,
                           const Eigen::MatrixBase<Derived2>& voxel_center,
                           const Eigen::ArrayBase<Derived3>& feat,
                           const size_t idx) {
          bool new_nearest_neighbor = false;
          TReal sqr_d = 0;
          if (POS_FN == NEAREST_NEIGHBOR || FEAT_FN == NEAREST_NEIGHBOR) {
              sqr_d = (voxel_center - pos).squaredNorm();
              if (sqr_d < min_sqr_dist_to_center_) {
                  new_nearest_neighbor = true;
                  min_sqr_dist_to_center_ = sqr_d;
              }
          }
  
          if (POS_FN == AVERAGE) {
              position_ += pos.array();
          } else if (POS_FN == NEAREST_NEIGHBOR && new_nearest_neighbor) {
              position_ = pos;
          } else if (POS_FN == CENTER) {
              if (count_ == 0) position_ = voxel_center;
          }
  
          if (count_ == 0) {
              features_.resizeLike(feat);
              features_.setZero();
              if (FEAT_FN == NEAREST_NEIGHBOR) {
                  features_ = feat;
                  index_.resize(1);
                  index_(0) = idx;
                  ++count_;
                  return;
              } else if (FEAT_FN == MAX) {
                  features_ = feat;
                  index_.resizeLike(feat);
                  index_ = idx;
                  ++count_;
                  return;
              }
          }
          if (FEAT_FN == AVERAGE) {
              features_ += feat;
          } else if (FEAT_FN == NEAREST_NEIGHBOR && new_nearest_neighbor) {
              features_ = feat;
              index_(0) = idx;
          } else if (FEAT_FN == MAX) {
              for (int i = 0; i < features_.rows(); ++i) {
                  if (feat(i) > features_(i)) {
                      features_(i) = feat(i);
                      index_(i) = idx;
                  }
              }
          }
          ++count_;
      }
  
      inline Eigen::Array<TReal, 3, 1> Position() const {
          if (POS_FN == AVERAGE) {
              return position_ / count_;
          } else  // if( POS_FN == NEAREST_NEIGHBOR || POS_FN == CENTER )
          {
              return position_;
          }
      }
  
      inline Eigen::Array<TFeat, Eigen::Dynamic, 1> Features() const {
          if (FEAT_FN == AVERAGE) {
              return features_ / count_;
          } else  // if( FEAT_FN == NEAREST_NEIGHBOR || FEAT_FN == MAX )
          {
              return features_;
          }
      }
  
      inline int Count() const { return count_; }
  
      inline Eigen::Array<size_t, Eigen::Dynamic, 1> Index() const {
          return index_;
      }
  
  private:
      int count_;
      TReal min_sqr_dist_to_center_;
      Eigen::Array<TReal, 3, 1> position_;
      Eigen::Array<TFeat, Eigen::Dynamic, 1> features_;
      Eigen::Array<size_t, Eigen::Dynamic, 1> index_;
  };
  
  /// Function for debugging. Checks if the voxel size is too small.
  template <class T>
  bool CheckVoxelSize(std::string& err,
                      const size_t num_positions,
                      const T* const positions,
                      const T voxel_size) {
      typedef Eigen::Array<double, 3, 1> Vec3_t;
      if (num_positions == 0) {
          return true;
      }
  
      Vec3_t bb_min, bb_max;
      bb_min << positions[0], positions[1], positions[2];
      bb_max = bb_min;
  
      Vec3_t voxel_size3(voxel_size, voxel_size, voxel_size);
  
      for (size_t i = 1; i < num_positions; ++i) {
          Vec3_t pos(positions[i * 3 + 0], positions[i * 3 + 1],
                     positions[i * 3 + 2]);
          bb_min = bb_min.min(pos);
          bb_max = bb_max.max(pos);
      }
  
      // the min and max bounding box shall be a multiple of the voxel size
      bb_min /= voxel_size3;
      bb_min = bb_min.floor() * voxel_size3;
      bb_max /= voxel_size3;
      bb_max = bb_max.ceil() * voxel_size3;
  
      if (voxel_size * double(std::numeric_limits<int>::max()) <
                  bb_max.maxCoeff() ||
          voxel_size * double(std::numeric_limits<int>::min()) >
                  bb_min.maxCoeff()) {
          err = "voxel_size is too small\n";
          return false;
      }
      return true;
  }
  
  /// Computes an integer voxel index for a 3D position.
  ///
  /// \param pos               A 3D position.
  /// \param inv_voxel_size    The reciprocal of the voxel size
  ///
  template <class TDerived>
  Eigen::Vector3i ComputeVoxelIndex(
          const Eigen::ArrayBase<TDerived>& pos,
          const typename TDerived::Scalar& inv_voxel_size) {
      typedef typename TDerived::Scalar Scalar_t;
      Eigen::Array<Scalar_t, 3, 1> ref_coord = pos * inv_voxel_size;
  
      Eigen::Vector3i voxel_index;
      voxel_index = ref_coord.floor().template cast<int>();
      return voxel_index;
  }
  
  // implementation for VoxelPooling with template parameter for the accumulator.
  template <class TReal, class TFeat, class ACCUMULATOR, class OUTPUT_ALLOCATOR>
  void _VoxelPooling(size_t num_inp,
                     const TReal* const inp_positions,
                     int in_channels,
                     const TFeat* inp_features,
                     TReal voxel_size,
                     OUTPUT_ALLOCATOR& output_allocator) {
      if (num_inp == 0) {
          TReal* out_pos_ptr;
          TFeat* out_feat_ptr;
          output_allocator.AllocPooledPositions(&out_pos_ptr, 0);
          output_allocator.AllocPooledFeatures(&out_feat_ptr, 0, in_channels);
          return;
      }
  
      typedef Eigen::Array<TReal, 3, 1> Vec3_t;
      typedef Eigen::Array<TFeat, Eigen::Dynamic, 1> FeatureVec_t;
  
      std::unordered_map<Eigen::Vector3i, ACCUMULATOR,
                         open3d::utility::hash_eigen::hash<Eigen::Vector3i>>
              voxelindex_to_accpoint;
  
      Vec3_t voxel_center;
      Eigen::Vector3i voxel_index;
      TReal inv_voxel_size = 1 / voxel_size;
      TReal half_voxel_size = 0.5 * voxel_size;
      for (size_t i = 0; i < num_inp; ++i) {
          Eigen::Map<const Vec3_t> pos(inp_positions + i * 3);
  
          voxel_index = ComputeVoxelIndex(pos, inv_voxel_size);
  
          voxel_center << voxel_index(0) * voxel_size + half_voxel_size,
                  voxel_index(1) * voxel_size + half_voxel_size,
                  voxel_index(2) * voxel_size + half_voxel_size;
  
          Eigen::Map<const FeatureVec_t> feat(inp_features + in_channels * i,
                                              in_channels);
          voxelindex_to_accpoint[voxel_index].AddPoint(
                  pos.matrix(), voxel_center.matrix(), feat);
      }
  
      const size_t num_out = voxelindex_to_accpoint.size();
  
      TReal* out_pos_ptr;
      TFeat* out_feat_ptr;
      output_allocator.AllocPooledPositions(&out_pos_ptr, num_out);
      output_allocator.AllocPooledFeatures(&out_feat_ptr, num_out, in_channels);
  
      size_t i = 0;
      for (const auto point : voxelindex_to_accpoint) {
          Vec3_t pos = point.second.Position();
          Eigen::Map<Vec3_t> out_pos(out_pos_ptr + i * 3);
          out_pos = pos;
  
          Eigen::Map<FeatureVec_t> out_feat(out_feat_ptr + i * in_channels,
                                            in_channels);
          out_feat = point.second.Features();
          ++i;
      }
  }
  
  // implementation for VoxelPoolingBackprop with template parameter for the
  // accumulator.
  template <class TReal, class TFeat, class ACCUMULATOR, AccumulationFn FEAT_FN>
  void _VoxelPoolingBackprop(TFeat* features_backprop,
                             size_t num_inp,
                             const TReal* const inp_positions,
                             int in_channels,
                             const TFeat* const inp_features,
                             size_t num_pooled,
                             const TReal* const pooled_positions,
                             const TFeat* const pooled_features_gradient,
                             TReal voxel_size) {
      if (num_inp == 0) {
          return;
      }
      memset(features_backprop, 0, sizeof(TFeat) * num_inp * in_channels);
  
      typedef Eigen::Array<TReal, 3, 1> Vec3_t;
      typedef Eigen::Array<TFeat, Eigen::Dynamic, 1> FeatureVec_t;
  
      Vec3_t voxel_size3(voxel_size, voxel_size, voxel_size);
  
      // create the two hash maps in parallel
      tbb::task_group task_group;
  
      std::unordered_map<Eigen::Vector3i, ACCUMULATOR,
                         open3d::utility::hash_eigen::hash<Eigen::Vector3i>>
              voxelindex_to_accpoint;
  
      task_group.run([&] {
          Vec3_t voxel_center;
          Eigen::Vector3i voxel_index;
          TReal inv_voxel_size = 1 / voxel_size;
          TReal half_voxel_size = 0.5 * voxel_size;
          for (size_t i = 0; i < num_inp; ++i) {
              Eigen::Map<const Vec3_t> pos(inp_positions + i * 3);
  
              voxel_index = ComputeVoxelIndex(pos, inv_voxel_size);
  
              voxel_center << voxel_index(0) * voxel_size + half_voxel_size,
                      voxel_index(1) * voxel_size + half_voxel_size,
                      voxel_index(2) * voxel_size + half_voxel_size;
  
              Eigen::Map<const FeatureVec_t> feat(inp_features + in_channels * i,
                                                  in_channels);
              voxelindex_to_accpoint[voxel_index].AddPoint(
                      pos.matrix(), voxel_center.matrix(), feat, i);
          }
      });
  
      std::unordered_map<Eigen::Vector3i, size_t,
                         open3d::utility::hash_eigen::hash<Eigen::Vector3i>>
              voxelindex_to_gradindex;
  
      task_group.run([&] {
          Eigen::Vector3i voxel_index;
          TReal inv_voxel_size = 1 / voxel_size;
          for (size_t i = 0; i < num_pooled; ++i) {
              Eigen::Map<const Vec3_t> pos(pooled_positions + i * 3);
  
              voxel_index = ComputeVoxelIndex(pos, inv_voxel_size);
  
              voxelindex_to_gradindex[voxel_index] = i;
          }
      });
  
      task_group.wait();
  
      if (FEAT_FN == AVERAGE) {
          Eigen::Vector3i voxel_index;
          TReal inv_voxel_size = 1 / voxel_size;
          for (size_t i = 0; i < num_inp; ++i) {
              Eigen::Map<const Vec3_t> pos(inp_positions + i * 3);
  
              voxel_index = ComputeVoxelIndex(pos, inv_voxel_size);
  
              Eigen::Map<FeatureVec_t> feat_bp(
                      features_backprop + in_channels * i, in_channels);
  
              size_t grad_idx = voxelindex_to_gradindex[voxel_index];
              int count = voxelindex_to_accpoint[voxel_index].Count();
              Eigen::Map<const FeatureVec_t> grad(
                      pooled_features_gradient + in_channels * grad_idx,
                      in_channels);
              feat_bp = grad / count;
          }
      }
  
      if (FEAT_FN == NEAREST_NEIGHBOR) {
          for (const auto point : voxelindex_to_accpoint) {
              size_t idx = point.second.Index()(0);
              Eigen::Map<FeatureVec_t> feat_bp(
                      features_backprop + in_channels * idx, in_channels);
  
              size_t grad_idx = voxelindex_to_gradindex[point.first];
              Eigen::Map<const FeatureVec_t> grad(
                      pooled_features_gradient + in_channels * grad_idx,
                      in_channels);
              feat_bp = grad;
          }
      }
  
      if (FEAT_FN == MAX) {
          for (const auto point : voxelindex_to_accpoint) {
              size_t grad_idx = voxelindex_to_gradindex[point.first];
              Eigen::Map<const FeatureVec_t> grad(
                      pooled_features_gradient + in_channels * grad_idx,
                      in_channels);
              for (int i = 0; i < in_channels; ++i) {
                  size_t idx = point.second.Index()(i);
                  Eigen::Map<FeatureVec_t> feat_bp(
                          features_backprop + in_channels * idx, in_channels);
                  feat_bp(i) = grad(i);
              }
          }
      }
  }
  
  }  // namespace
  
  /// Pooling operation for point clouds. Aggregates points that are inside the
  /// same voxel.
  ///
  /// \tparam TReal    Scalar type for point positions.
  ///
  /// \tparam TFeat    Scalar type for the features.
  ///
  /// \tparam OUTPUT_ALLOCATOR    Type of the output_allocator. See
  ///         \p output_allocator for more information.
  ///
  /// \param num_inp    The number of input points.
  ///
  /// \param inp_positions    Array with 3D point positions.
  ///
  /// \param in_channels    The number of input feature channels.
  ///
  /// \param inp_features    The array with the point features. The shape is
  ///                        [num_inp, in_channels].
  ///
  /// \param voxel_size    The voxel size (voxel edge length) used for pooling.
  ///
  /// \param output_allocator    An object that implements functions for
  ///         allocating the output arrays. The object must implement functions
  ///         AllocPooledPositions(TReal** ptr, size_t num) and
  ///         AllocPooledFeatures(TFeat** ptr, size_t num, channels), with 'num'
  ///         as the number of output points and 'channels' as \p in_channels .
  ///         Both functions should allocate memory and return a pointer
  ///         to that memory in ptr.
  ///         Both functions must accept the argument num==0. In this case ptr
  ///         does not need to be set. Note that AllocPooledPositions must
  ///         allocate memory for num*3*sizeof(TReal) bytes.
  ///
  /// \param position_fn    Defines how the new point positions will be computed.
  ///        AVERAGE computes the center of gravity for the points within one
  ///        voxel.
  ///        NEAREST_NEIGHBOR selects the point closest to the voxel center.
  ///        CENTER uses the voxel center for the position of the generated point.
  ///
  /// \param feature_fn    Defines how the new features will be computed.
  ///        AVERAGE computes the average feature vector.
  ///        NEAREST_NEIGHBOR selects the feature vector of the point closest to
  ///        the voxel center.
  ///        MAX uses the maximum feature among all points within the voxel.
  ///
  template <class TReal, class TFeat, class OUTPUT_ALLOCATOR>
  void VoxelPooling(size_t num_inp,
                    const TReal* const inp_positions,
                    int in_channels,
                    const TFeat* inp_features,
                    TReal voxel_size,
                    OUTPUT_ALLOCATOR& output_allocator,
                    AccumulationFn position_fn,
                    AccumulationFn feature_fn) {
  #define CALL_TEMPLATE(POS_FN, FEAT_FN)                                         \
      if (POS_FN == position_fn && FEAT_FN == feature_fn) {                      \
          _VoxelPooling<TReal, TFeat,                                            \
                        Accumulator<TReal, TFeat, POS_FN, FEAT_FN>>(             \
                  num_inp, inp_positions, in_channels, inp_features, voxel_size, \
                  output_allocator);                                             \
      }
  
      CALL_TEMPLATE(AVERAGE, AVERAGE)
      CALL_TEMPLATE(AVERAGE, NEAREST_NEIGHBOR)
      CALL_TEMPLATE(AVERAGE, MAX)
      CALL_TEMPLATE(NEAREST_NEIGHBOR, AVERAGE)
      CALL_TEMPLATE(NEAREST_NEIGHBOR, NEAREST_NEIGHBOR)
      CALL_TEMPLATE(NEAREST_NEIGHBOR, MAX)
      CALL_TEMPLATE(CENTER, AVERAGE)
      CALL_TEMPLATE(CENTER, NEAREST_NEIGHBOR)
      CALL_TEMPLATE(CENTER, MAX)
  
  #undef CALL_TEMPLATE
  }
  
  /// Backpropagation to the features for VoxelPooling.
  ///
  /// \param features_backprop    The output array with the gradients for the
  ///        features.
  ///
  /// \param num_pooled    The number of points after pooling with \p
  ///         VoxelPooling.
  ///
  /// \param pooled_positions    Array with the 3D positions after pooling.
  ///
  /// \param pooled_features_gradient    Array with the point features after
  ///        pooling.
  ///
  /// See \p VoxelPooling for the description of the remaining parameters.
  ///
  template <class TReal, class TFeat>
  void VoxelPoolingBackprop(TFeat* features_backprop,
                            size_t num_inp,
                            const TReal* const inp_positions,
                            int in_channels,
                            const TFeat* const inp_features,
                            size_t num_pooled,
                            const TReal* const pooled_positions,
                            const TFeat* const pooled_features_gradient,
                            TReal voxel_size,
                            AccumulationFn position_fn,
                            AccumulationFn feature_fn) {
  #define CALL_TEMPLATE(POS_FN, FEAT_FN)                                        \
      if (POS_FN == position_fn && FEAT_FN == feature_fn) {                     \
          _VoxelPoolingBackprop<                                                \
                  TReal, TFeat,                                                 \
                  AccumulatorBackprop<TReal, TFeat, POS_FN, FEAT_FN>, FEAT_FN>( \
                  features_backprop, num_inp, inp_positions, in_channels,       \
                  inp_features, num_pooled, pooled_positions,                   \
                  pooled_features_gradient, voxel_size);                        \
      }
  
      CALL_TEMPLATE(AVERAGE, AVERAGE)
      CALL_TEMPLATE(AVERAGE, NEAREST_NEIGHBOR)
      CALL_TEMPLATE(AVERAGE, MAX)
      CALL_TEMPLATE(NEAREST_NEIGHBOR, AVERAGE)
      CALL_TEMPLATE(NEAREST_NEIGHBOR, NEAREST_NEIGHBOR)
      CALL_TEMPLATE(NEAREST_NEIGHBOR, MAX)
      CALL_TEMPLATE(CENTER, AVERAGE)
      CALL_TEMPLATE(CENTER, NEAREST_NEIGHBOR)
      CALL_TEMPLATE(CENTER, MAX)
  
  #undef CALL_TEMPLATE
  }
  
  }  // namespace detail
  }  // namespace ml
  }  // namespace open3d
  