// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2023 www.open3d.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------

#pragma once

#include <tbb/parallel_for.h>

#include "open3d/ml/impl/continuous_conv/CoordinateTransformation.h"

namespace open3d {
namespace ml {
namespace impl {

/// Implementation of CConvComputeFeatures with template parameters for
/// configuration.
template <class TFeat,
          class TOut,
          class TReal,
          class TIndex,
          InterpolationMode INTERPOLATION,
          CoordinateMapping MAPPING,
          bool ALIGN_CORNERS,
          bool INDIVIDUAL_EXTENT,
          bool ISOTROPIC_EXTENT,
          bool POINT_IMPORTANCE>
void _CConvComputeFeaturesCPU(TOut* out_features,
                              const std::vector<int>& filter_dims,
                              const TFeat* filter,
                              size_t num_out,
                              const TReal* out_positions,
                              size_t num_inp,
                              const TReal* inp_positions,
                              const TFeat* inp_features,
                              const TFeat* inp_importance,
                              size_t neighbors_index_size,
                              const TIndex* neighbors_index,
                              const TFeat* neighbors_importance,
                              const int64_t* neighbors_row_splits,
                              const TReal* extents,
                              const TReal* offsets,
                              bool normalize) {
    const bool NEIGHBORS_IMPORTANCE = neighbors_importance != nullptr;
    const int VECSIZE = 32;
    typedef Eigen::Array<TReal, VECSIZE, 1> Vec_t;
    typedef InterpolationVec<TReal, VECSIZE, INTERPOLATION> InterpolationVec_t;
    InterpolationVec_t interpolation;

    const int in_channels = filter_dims[filter_dims.size() - 2];
    const int out_channels = filter_dims[filter_dims.size() - 1];

    int spatial_filter_size = 1;
    for (int i = 0; i < 3; ++i) spatial_filter_size *= filter_dims[i];
    Eigen::Array<int, 3, 1> filter_size_xyz(filter_dims[2], filter_dims[1],
                                            filter_dims[0]);

    memset(out_features, 0, sizeof(TOut) * num_out * out_channels);

    tbb::parallel_for(
            tbb::blocked_range<size_t>(0, num_out, 32),
            [&](const tbb::blocked_range<size_t>& r) {
                int range_length = r.end() - r.begin();

                Eigen::Matrix<TOut, Eigen::Dynamic, 1> normalizers(range_length,
                                                                   1);
                normalizers.setZero();

                Eigen::Matrix<TFeat, Eigen::Dynamic, Eigen::Dynamic> B(
                        in_channels * spatial_filter_size, range_length);
                B.setZero();

                typedef Eigen::Array<TFeat, VECSIZE, Eigen::Dynamic> Matrix;
                Matrix infeat(VECSIZE, in_channels);

                Eigen::Array<TReal, 3, 1> offsets_(offsets[0], offsets[1],
                                                   offsets[2]);

                Eigen::Array<TReal, VECSIZE, 3> inv_extents;
                if (INDIVIDUAL_EXTENT == false) {
                    if (ISOTROPIC_EXTENT) {
                        inv_extents = 1 / extents[0];
                    } else {
                        inv_extents.col(0) = 1 / extents[0];
                        inv_extents.col(1) = 1 / extents[1];
                        inv_extents.col(2) = 1 / extents[2];
                    }
                }

                for (size_t out_idx = r.begin(); out_idx != r.end();
                     ++out_idx) {
                    const int out_col = out_idx - r.begin();
                    const size_t neighbor_start = neighbors_row_splits[out_idx];
                    const size_t neighbor_end =
                            neighbors_row_splits[out_idx + 1];

                    if (INDIVIDUAL_EXTENT) {
                        if (ISOTROPIC_EXTENT) {
                            inv_extents = 1 / extents[out_idx];
                        } else {
                            inv_extents.col(0) = 1 / extents[3 * out_idx + 0];
                            inv_extents.col(1) = 1 / extents[3 * out_idx + 1];
                            inv_extents.col(2) = 1 / extents[3 * out_idx + 2];
                        }
                    }

                    typename InterpolationVec_t::Weight_t interp_weights;
                    typename InterpolationVec_t::Idx_t interp_indices;

                    int vec_valid_count = 0;
                    Vec_t x, y, z;

                    // set to zero to avoid problems with vectors with less than
                    // VECSIZE valid entries
                    x.setZero();
                    y.setZero();
                    z.setZero();
                    for (size_t n = neighbor_start; n < neighbor_end; ++n) {
                        const size_t inp_idx = neighbors_index[n];
                        const int i = vec_valid_count;
                        x(i) = inp_positions[inp_idx * 3 + 0] -
                               out_positions[out_idx * 3 + 0];
                        y(i) = inp_positions[inp_idx * 3 + 1] -
                               out_positions[out_idx * 3 + 1];
                        z(i) = inp_positions[inp_idx * 3 + 2] -
                               out_positions[out_idx * 3 + 2];

                        const TFeat n_importance =
                                (NEIGHBORS_IMPORTANCE ? neighbors_importance[n]
                                                      : TFeat(1));
                        normalizers(out_col) += TOut(n_importance);

                        for (int ic = 0; ic < in_channels; ++ic)
                            infeat(i, ic) =
                                    inp_features[inp_idx * in_channels + ic];

                        TFeat importance(1.0);
                        if (POINT_IMPORTANCE)
                            importance = inp_importance[inp_idx];
                        if (NEIGHBORS_IMPORTANCE) importance *= n_importance;

                        if (POINT_IMPORTANCE || NEIGHBORS_IMPORTANCE) {
                            for (int ic = 0; ic < in_channels; ++ic)
                                infeat(i, ic) *= importance;
                        }

                        ++vec_valid_count;
                        if (vec_valid_count == VECSIZE) {
                            ComputeFilterCoordinates<ALIGN_CORNERS, MAPPING>(
                                    x, y, z, filter_size_xyz, inv_extents,
                                    offsets_);
                            interpolation.Interpolate(
                                    interp_weights, interp_indices, x, y, z,
                                    filter_size_xyz, in_channels);
                            for (int k = 0; k < VECSIZE; ++k)
                                for (int j = 0; j < InterpolationVec_t::Size();
                                     ++j) {
                                    for (int ic = 0; ic < in_channels; ++ic)
                                        B(interp_indices(j, k) + ic, out_col) +=
                                                TFeat(interp_weights(j, k)) *
                                                infeat(k, ic);
                                }
                            vec_valid_count = 0;
                        }
                    }
                    if (vec_valid_count) {
                        ComputeFilterCoordinates<ALIGN_CORNERS, MAPPING>(
                                x, y, z, filter_size_xyz, inv_extents,
                                offsets_);
                        interpolation.Interpolate(interp_weights,
                                                  interp_indices, x, y, z,
                                                  filter_size_xyz, in_channels);
                        for (int k = 0; k < vec_valid_count; ++k)
                            for (int j = 0; j < InterpolationVec_t::Size();
                                 ++j) {
                                for (int ic = 0; ic < in_channels; ++ic)
                                    B(interp_indices(j, k) + ic, out_col) +=
                                            TFeat(interp_weights(j, k)) *
                                            infeat(k, ic);
                            }
                    }

                }  // out_idx

                Eigen::Map<const Eigen::Matrix<TFeat, Eigen::Dynamic,
                                               Eigen::Dynamic>>
                        A(filter, out_channels,
                          spatial_filter_size * in_channels);
                Eigen::Map<Eigen::Matrix<TOut, Eigen::Dynamic, Eigen::Dynamic>>
                        C(out_features + (r.begin() * out_channels),
                          out_channels, range_length);

                C = (A * B).template cast<TOut>();
                if (normalize) {
                    for (int i = 0; i < range_length; ++i) {
                        if (normalizers(i) != TOut(0))
                            C.col(i) /= normalizers(i);
                    }
                }
            });
}

/// Computes the output features of a continuous convolution.
///
/// \tparam TFeat    Type for the features and weights
/// \tparam TOut     Type for the output features
/// \tparam TReal    Type for point positions and extents
/// \tparam TIndex   Type for neighbor indexing
///
/// \param out_features    Output array for the computed features with shape
///        [num_out, out channels]
///
/// \param filter_dims    The sizes of the filter dimensions. The size of
///        filter_dims must be 5. The order is
///        [depth, height, width, inp channels, out channels].
///
/// \param filter    Pointer to the filter values.
///
/// \param num_out    The number of output points.
///
/// \param out_positions    The positions of the output points. The shape is
///        [num_out, 3].
///
/// \param num_inp    The number of input points.
///
/// \param inp_positions    The positions of the input points. The shape is
///        [num_inp, 3].
///
/// \param inp_features    The input features with shape
///        [num_inp, in_channels].
///
/// \param inp_importance    Optional importance for each input point with
///        shape [num_inp]. Set to null to disable.
///
/// \param neighbors_index_size    The size of the neighbors_index array.
///
/// \param neighbors_index    The array with lists of neighbors for each
///        output point. The start and end of each sublist is defined by
///        \p neighbors_row_splits.
///
/// \param neighbors_importance    Optional importance for each entry in
///        \p neighbors_index. Set to null to disable.
///
/// \param neighbors_row_splits   The prefix sum which defines the start
///        and end of the sublists in \p neighbors_index. The size of the
///        array is \p num_out + 1.
///
/// \param extents    The spatial extents of the filter in coordinate units.
///        extents can be a scalar or a 1D array of shape [num_out] or a
///        2D array of shape [num_out,3]. The shape depends on
///        \p individual_extent and \p isotropic_extent.
///
/// \param offsets    A single 3D vector used in the filter coordinate
///        computation. The shape is [3].
///
/// \param interpolation    The interpolation mode. Either LINEAR or
///        NEAREST_NEIGHBOR.
///
/// \param coordinate_mapping    The coordinate mapping function. One of
///        IDENTITY, BALL_TO_CUBE_RADIAL, BALL_TO_CUBE_VOLUME_PRESERVING.
///
/// \param align_corners    If true then the voxel centers of the outer voxels
///        of the filter array are mapped to the boundary of the filter shape.
///        If false then the boundary of the filter array is mapped to the
///        boundary of the filter shape.
///
/// \param individual_extent    If true each output point has an individual
///        extent.
///
/// \param isotropic_extent    If true each then the extent is isotropic for
///        each output point.
///
/// \param normalize    If true then the result is normalized either by the
///        number of points (neighbors_importance is null) or by the sum of
///        the respective values in neighbors_importance.
///
template <class TFeat, class TOut, class TReal, class TIndex>
void CConvComputeFeaturesCPU(TOut* out_features,
                             const std::vector<int>& filter_dims,
                             const TFeat* filter,
                             size_t num_out,
                             const TReal* out_positions,
                             size_t num_inp,
                             const TReal* inp_positions,
                             const TFeat* inp_features,
                             const TFeat* inp_importance,
                             size_t neighbors_index_size,
                             const TIndex* neighbors_index,
                             const TFeat* neighbors_importance,
                             const int64_t* neighbors_row_splits,
                             const TReal* extents,
                             const TReal* offsets,
                             InterpolationMode interpolation,
                             CoordinateMapping coordinate_mapping,
                             bool align_corners,
                             bool individual_extent,
                             bool isotropic_extent,
                             bool normalize) {
    // Dispatch all template parameter combinations
    bool has_importance = inp_importance;

#define FN_PARAMETERS                                                          \
    out_features, filter_dims, filter, num_out, out_positions, num_inp,        \
            inp_positions, inp_features, inp_importance, neighbors_index_size, \
            neighbors_index, neighbors_importance, neighbors_row_splits,       \
            extents, offsets, normalize

#define CALL_TEMPLATE(INTERPOLATION, MAPPING, ALIGN_CORNERS,                \
                      INDIVIDUAL_EXTENT, ISOTROPIC_EXTENT, HAS_IMPORTANCE)  \
    if (INTERPOLATION == interpolation && MAPPING == coordinate_mapping &&  \
        ALIGN_CORNERS == align_corners &&                                   \
        INDIVIDUAL_EXTENT == individual_extent &&                           \
        ISOTROPIC_EXTENT == isotropic_extent &&                             \
        HAS_IMPORTANCE == has_importance)                                   \
        _CConvComputeFeaturesCPU<TFeat, TOut, TReal, TIndex, INTERPOLATION, \
                                 MAPPING, ALIGN_CORNERS, INDIVIDUAL_EXTENT, \
                                 ISOTROPIC_EXTENT, HAS_IMPORTANCE>(         \
                FN_PARAMETERS);

#define CALL_TEMPLATE2(INTERPOLATION, MAPPING)                       \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, true, true, true)    \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, true, true, false)   \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, true, false, true)   \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, true, false, false)  \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, false, true, true)   \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, false, true, false)  \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, false, false, true)  \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, false, false, false) \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, true, true, true)   \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, true, true, false)  \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, true, false, true)  \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, true, false, false) \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, false, true, true)  \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, false, true, false) \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, false, false, true) \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, false, false, false)

#define CALL_TEMPLATE3(INTERPOLATION)                                     \
    CALL_TEMPLATE2(INTERPOLATION, CoordinateMapping::BALL_TO_CUBE_RADIAL) \
    CALL_TEMPLATE2(INTERPOLATION,                                         \
                   CoordinateMapping::BALL_TO_CUBE_VOLUME_PRESERVING)     \
    CALL_TEMPLATE2(INTERPOLATION, CoordinateMapping::IDENTITY)

#define CALL_TEMPLATE4                               \
    CALL_TEMPLATE3(InterpolationMode::LINEAR)        \
    CALL_TEMPLATE3(InterpolationMode::LINEAR_BORDER) \
    CALL_TEMPLATE3(InterpolationMode::NEAREST_NEIGHBOR)

    CALL_TEMPLATE4

#undef CALL_TEMPLATE
#undef CALL_TEMPLATE2
#undef CALL_TEMPLATE3
#undef CALL_TEMPLATE4

#undef FN_PARAMETERS
}

}  // namespace impl
}  // namespace ml
}  // namespace open3d