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

#include "open3d/ml/tensorflow/TensorFlowHelper.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"

using namespace tensorflow;

REGISTER_OP("Open3DTrilinearDevoxelize")
        .Attr("resolution: int")
        .Attr("is_training: bool")
        .Input("coords: float32")
        .Input("features: float32")
        .Output("outputs: float32")
        .Output("indices: int32")
        .Output("weights: float32")
        .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
            using namespace ::tensorflow::shape_inference;
            using namespace open3d::ml::op_util;
            ShapeHandle coords;    // (batch_size, 3, N)
            ShapeHandle features;  // (batch_size, C, R, R, R)

            TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 3, &coords));
            TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 5, &features));

            Dim batch_size("batch_size");
            Dim num_points("num_points");
            Dim resolution("resolution");
            Dim feat_dim("feat_dim");
            CHECK_SHAPE_HANDLE(c, coords, batch_size, 3, num_points);
            CHECK_SHAPE_HANDLE(c, features, batch_size, feat_dim, resolution,
                               resolution, resolution);

            ShapeHandle out1, out2;
            out1 = c->MakeShape(
                    {c->MakeDim(batch_size.value()),
                     c->MakeDim(feat_dim.value()),
                     c->MakeDim(num_points.value())});  // (batch_size, C, N)
            out2 = c->MakeShape(
                    {c->MakeDim(batch_size.value()), 8,
                     c->MakeDim(num_points.value())});  // (batch_size, 8, N)

            c->set_output(0, out1);
            c->set_output(1, out2);
            c->set_output(2, out2);

            return Status();
        })
        .Doc(R"doc(
Trilinear Devoxelize.

This function takes a 3D voxel grid and a list of coordinates and
computes interpolated features corresponding to each point.

Minimal example::
  import open3d.ml.tf as ml3d

  coords = tf.Tensor(
        [[[0.2 0.0 0.0 1.0 1.5]
          [1.0 1.2 0.9 0.0 0.7]
          [0.2 0.6 0.8 0.7 1.1]]], shape=(1, 3, 5), dtype=float32)

  features = tf.Tensor(
  [[[[[0.  0.5]
      [0.6 0.4]]

     [[0.5 0.7]
      [0.6 0.5]]]

    [[[0.4 0.8]
      [0.6 0.3]]

     [[0.4 0.2]
      [0.8 0.6]]]

    [[[0.1 0.2]
      [0.  0.6]]

     [[0.9 0. ]
      [0.2 0.3]]]]], shape=(1, 3, 2, 2, 2), dtype=float32)

  ml3d.ops.trilinear_devoxelize(coords,
                                features,
                                resolution,
                                is_training)

  # returns output tf.Tensor(
  #                array([[[0.564     , 0.508     , 0.436     , 0.64      , 0.5005    ],
  #                        [0.58400005, 0.39200002, 0.396     , 0.26000002, 0.47900003],
  #                        [0.14      , 0.36      , 0.45000002, 0.27      , 0.0975    ]]],
  #                shape=(1, 3, 5), dtype=float32)
  #
  #         indices tf.Tensor([[[ 2,  2,  0,  4,  5],
  #                             [ 3,  3,  1,  5,  6],
  #                             [ 2,  4,  2,  4,  7],
  #                             [ 3,  5,  3,  5,  8],
  #                             [ 6,  2,  0,  4,  9],
  #                             [ 7,  3,  1,  5, 10],
  #                             [ 6,  4,  2,  4, 11],
  #                             [ 7,  5,  3,  5, 12]]],
  #                 shape=(1, 8, 5), dtype=float32)
  #
  #         weights tf.Tensor([[[0.64000005, 0.31999996, 0.02      , 0.3       , 0.135     ],
  #                             [0.16000001, 0.48      , 0.08000002, 0.7       , 0.015     ],
  #                             [0.        , 0.08000001, 0.17999998, 0.        , 0.315     ],
  #                             [0.        , 0.12000003, 0.71999997, 0.        , 0.03500001],
  #                             [0.16000001, 0.        , 0.        , 0.        , 0.135     ],
  #                             [0.04      , 0.        , 0.        , 0.        , 0.015     ],
  #                             [0.        , 0.        , 0.        , 0.        , 0.315     ],
  #                             [0.        , 0.        , 0.        , 0.        , 0.03500001]]],
  #                 shape=(1, 8, 5), dtype=float32)

coords: List of 3D coordinates for which features to be interpolated.
  The shape of this tensor is [B, 3, N]. The range of coordinates is
  [0, resolution-1]. If all of the adjacent position of any coordinate are out
  of range, then the interpolated features will be 0. Voxel centers are at integral
  values of voxel grid.

features: A voxel grid of shape [B, C, R, R, R]. Here R is resolution.

resolution: Integer attribute defining resolution of the voxel grid.

is_training: Boolean variable for training phase.

outputs: Features for each point. The shape of this tensor is [B, C, N].

indices: Indices which are used to interpolate features. Shape is [B, 8, N].

weights: Weights for each index used to interpolate features. Shape is [B, 8, N].

)doc");

REGISTER_OP("Open3DTrilinearDevoxelizeGrad")
        .Input("grad_y: float32")
        .Input("indices: int32")
        .Input("weights: float32")
        .Attr("resolution: int")
        .Output("grad_x: float32")
        .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
            using namespace ::tensorflow::shape_inference;
            using namespace open3d::ml::op_util;
            ShapeHandle grad_y;   // (batch_size, C, N)
            ShapeHandle indices;  // (batch_size, 8, N)

            TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 3, &grad_y));
            TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 3, &indices));

            Dim batch_size("batch_size");
            Dim feat_dim("feat_dim");
            Dim num_points("num_points");
            CHECK_SHAPE_HANDLE(c, grad_y, batch_size, feat_dim, num_points);
            CHECK_SHAPE_HANDLE(c, indices, batch_size, 8, num_points);

            ShapeHandle out;
            int r_;
            TF_RETURN_IF_ERROR(c->GetAttr("resolution", &r_));
            out = c->MakeShape({c->MakeDim(batch_size.value()),
                                c->MakeDim(feat_dim.value()), r_, r_,
                                r_});  // (batch_size, C, R, R, R)

            c->set_output(0, out);

            return Status();
        })
        .Doc(R"doc(
Gradient function for Trilinear Devoxelize op.

This function takes feature gradients and indices, weights returned from
the op and computes gradient for voxelgrid.

grad_y: Gradients for the interpolated features. Shape is [B, C, N].

indices: Indices which are used to interpolate features. Shape is [B, 8, N].

weights: Weights for each index used to interpolate features. Shape is [B, 8, N].

resolution: Integer attribute defining resolution of the grid.

grad_x: Output gradients for the voxel grid. Shape is [B, C, R, R, R]

)doc");
