// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018-2021 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.
// ----------------------------------------------------------------------------

#include <tensorflow/core/framework/op.h>
#include <tensorflow/core/framework/op_kernel.h>
#include <tensorflow/core/framework/shape_inference.h>
#include <tensorflow/core/lib/core/errors.h>

#include "open3d/ml/tensorflow/TensorFlowHelper.h"

using namespace tensorflow;

REGISTER_OP("Open3DSparseConvTransposeBackpropFilter")
        .Attr("TFeat: {float, double, bfloat16}")  // Type for features and
                                                   // weights
        .Attr("output_type: {float, double, bfloat16} = DT_FLOAT")  // Type for
                                                                    // the
                                                                    // output
                                                                    // features
        .Attr("TIndex: {int32, int64}")
        .Attr("TKernelIndex: {uint8, int16}")
        .Attr("normalize: bool = false")
        .Attr("max_temp_mem_MB: int = 64")
        .Input("filters: TFeat")  // [depth, height, width, in_ch, out_ch]
        .Input("out_importance: TFeat")                // [num_points_out]
        .Input("inp_features: TFeat")                  // [num_points_in, in_ch]
        .Input("inp_neighbors_importance_sum: TFeat")  // [num_points_in]
        .Input("inp_neighbors_row_splits: int64")      // [num_points_in]
        .Input("neighbors_index: TIndex")              // [?]
        .Input("neighbors_kernel_index: TKernelIndex")  // [?]
        .Input("neighbors_importance: TFeat")           // [?]
        .Input("neighbors_row_splits: int64")           // [num_points_out]
        .Input("out_features_gradient: TFeat")    // [num_points_out, out_ch]
        .Output("filter_backprop : output_type")  // [depth, height, width,
                                                  // in_ch, out_ch]
        .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
            using namespace ::tensorflow::shape_inference;
            using namespace open3d::ml::op_util;
            ShapeHandle filters = c->input(0);
            ShapeHandle out_importance = c->input(1);
            ShapeHandle inp_features = c->input(2);
            ShapeHandle inp_neighbors_importance_sum = c->input(3);
            ShapeHandle inp_neighbors_row_splits = c->input(4);
            ShapeHandle neighbors_index = c->input(5);
            ShapeHandle neighbors_kernel_index = c->input(6);
            ShapeHandle neighbors_importance = c->input(7);
            ShapeHandle neighbors_row_splits = c->input(8);
            ShapeHandle out_features_gradient = c->input(9);

            Dim num_out("num_out");
            Dim num_inp("num_inp");
            Dim num_kernel_elements("num_kernel_elements");
            Dim in_channels("in_channels");
            Dim out_channels("out_channels");
            Dim num_neighbors("num_neighbors");

            CHECK_SHAPE_HANDLE_COMBINE_FIRST_DIMS(
                    c, filters, num_kernel_elements, in_channels, out_channels);
            CHECK_SHAPE_HANDLE(c, neighbors_row_splits, num_out + 1);
            CHECK_SHAPE_HANDLE(c, out_importance, 0 || num_out);
            CHECK_SHAPE_HANDLE(c, inp_features, num_inp, in_channels);
            CHECK_SHAPE_HANDLE(c, inp_neighbors_importance_sum, 0 || num_inp);
            CHECK_SHAPE_HANDLE(c, inp_neighbors_row_splits, num_inp + 1);
            CHECK_SHAPE_HANDLE(c, neighbors_index, num_neighbors);
            CHECK_SHAPE_HANDLE(c, neighbors_kernel_index, num_neighbors);
            CHECK_SHAPE_HANDLE(c, neighbors_importance, 0 || num_neighbors);
            CHECK_SHAPE_HANDLE(c, out_features_gradient, num_out, out_channels);

            c->set_output(0, filters);
            return Status::OK();
        })
        .Doc(R"doc(
Computes the backrop for the filter of the SparseConvTranspose

normalize:
  If True the input feature values will be normalized by the number of neighbors.


max_temp_mem_MB:
  Defines the maximum temporary memory in megabytes to be used for the GPU
  implementation. More memory means fewer kernel invocations. Note that the
  a minimum amount of temp memory will always be allocated even if this
  variable is set to 0.


filters:
  The filter parameters.
  The shape of the filter is [depth, height, width, in_ch, out_ch].
  The dimensions 'depth', 'height', 'width' define the spatial resolution of
  the filter. The spatial size of the filter is defined by the parameter
  'extents'.


inp_features:
  A 2D tensor which stores a feature vector for each input point.


inp_neighbors_row_splits:
  The number of neighbors for each input point as exclusive prefix sum.


neighbors_index:
  The neighbors_index stores a list of indices of neighbors for each output point as nested lists.
  The start and end of each list can be computed using 'neighbors_row_splits'.


neighbors_kernel_index:
  Defines which kernel element to use for each neighbor. This array has the same length as neighbors_index.


neighbors_row_splits:
  The number of neighbors for each output point as exclusive prefix sum.


out_features_gradient:
  A Tensor with the gradient for the outputs of the SparseConvTranspose in the forward pass.

output_type: The type for the output.

filter_backprop:
  The gradients for the filter

)doc");
