/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"

namespace tensorflow {

using shape_inference::DimensionHandle;
using shape_inference::InferenceContext;
using shape_inference::ShapeHandle;

REGISTER_OP("SetSize")
    .Input("set_indices: int64")
    .Input("set_values: T")
    .Input("set_shape: int64")
    .Attr("validate_indices: bool = true")
    .Attr("T: {int8, int16, int32, int64, uint8, uint16, string}")
    .Output("size: int32")
    .SetShapeFn(shape_inference::UnknownShape);

REGISTER_OP("DenseToDenseSetOperation")
    .Input("set1: T")
    .Input("set2: T")
    .Attr("set_operation: string")
    .Attr("validate_indices: bool = true")
    .Attr("T: {int8, int16, int32, int64, uint8, uint16, string}")
    .Output("result_indices: int64")
    .Output("result_values: T")
    .Output("result_shape: int64")
    .SetShapeFn([](InferenceContext* c) {
      if (c->num_inputs() != 2) {
        return errors::InvalidArgument("len(inputs) != 2.");
      }
      // The following should stay in sync with `ComputeDenseToDense` shape
      // assertions in kernels/set_kernels.cc.
      // Dimension n contains the set values to be compared, so ranks must be
      // >= 2, and the first n-1 dimensions of inputs and output must be
      // compatible.
      DimensionHandle output_rank;
      ShapeHandle input0_shape = c->input(0);
      TF_RETURN_IF_ERROR(c->WithRankAtLeast(input0_shape, 2, &input0_shape));
      if (c->RankKnown(input0_shape)) {
        const int32 input0_rank = c->Rank(input0_shape);
        ShapeHandle input1_shape = c->input(1);
        TF_RETURN_IF_ERROR(
            c->WithRank(input1_shape, input0_rank, &input1_shape));
        if (c->RankKnown(input1_shape)) {
          // If both ranks are specified, the first n-1 dims must be compatible.
          const int32 rank = c->Rank(input1_shape);
          ShapeHandle group0_shape;
          TF_RETURN_IF_ERROR(
              c->Subshape(input0_shape, 0, rank - 1, &group0_shape));
          ShapeHandle group1_shape;
          TF_RETURN_IF_ERROR(
              c->Subshape(input1_shape, 0, rank - 1, &group1_shape));
          ShapeHandle unused_shape;
          TF_RETURN_IF_ERROR(
              c->Merge(group0_shape, group1_shape, &unused_shape));
        }
        output_rank = c->MakeDim(input0_rank);
      } else {
        ShapeHandle input1_shape = c->input(1);
        TF_RETURN_IF_ERROR(c->WithRankAtLeast(input1_shape, 2, &input1_shape));
        if (c->RankKnown(input1_shape)) {
          output_rank = c->MakeDim(c->Rank(input1_shape));
        } else {
          output_rank = c->UnknownDim();
        }
      }

      c->set_output(0, c->Matrix(c->UnknownDim(), output_rank));
      c->set_output(1, c->Vector(c->UnknownDim()));
      c->set_output(2, c->Vector(output_rank));
      return Status::OK();
    });

REGISTER_OP("DenseToSparseSetOperation")
    .Input("set1: T")
    .Input("set2_indices: int64")
    .Input("set2_values: T")
    .Input("set2_shape: int64")
    .Attr("set_operation: string")
    .Attr("validate_indices: bool = true")
    .Attr("T: {int8, int16, int32, int64, uint8, uint16, string}")
    .Output("result_indices: int64")
    .Output("result_values: T")
    .Output("result_shape: int64")
    .SetShapeFn([](InferenceContext* c) {
      if (c->num_inputs() != 4) {
        return errors::InvalidArgument("len(inputs) != 4.");
      }
      // The following should stay in sync with `ComputeDenseToSparse` shape
      // assertions in kernels/set_kernels.cc.
      // Ranks must be compatible, and be >= 2.
      ShapeHandle input1_shape_shape = c->input(3);
      TF_RETURN_IF_ERROR(shape_inference::ValidateSparseTensor(
          c, c->input(1), c->input(2), input1_shape_shape));

      DimensionHandle input1_rank_dim = c->Dim(input1_shape_shape, 0);

      DimensionHandle output_rank_dim;
      ShapeHandle input0_shape = c->input(0);
      TF_RETURN_IF_ERROR(c->WithRankAtLeast(input0_shape, 2, &input0_shape));
      if (c->RankKnown(input0_shape)) {
        const int32 input0_rank = c->Rank(input0_shape);
        TF_RETURN_IF_ERROR(
            c->WithValue(input1_rank_dim, input0_rank, &input1_rank_dim));
        output_rank_dim = c->MakeDim(input0_rank);
      } else if (c->ValueKnown(input1_rank_dim)) {
        output_rank_dim = input1_rank_dim;
      } else {
        output_rank_dim = c->UnknownDim();
      }

      c->set_output(0, c->Matrix(c->UnknownDim(), output_rank_dim));
      c->set_output(1, c->Vector(c->UnknownDim()));
      c->set_output(2, c->Vector(output_rank_dim));
      return Status::OK();
    });

REGISTER_OP("SparseToSparseSetOperation")
    .Input("set1_indices: int64")
    .Input("set1_values: T")
    .Input("set1_shape: int64")
    .Input("set2_indices: int64")
    .Input("set2_values: T")
    .Input("set2_shape: int64")
    .Attr("set_operation: string")
    .Attr("validate_indices: bool = true")
    .Attr("T: {int8, int16, int32, int64, uint8, uint16, string}")
    .Output("result_indices: int64")
    .Output("result_values: T")
    .Output("result_shape: int64")
    .SetShapeFn([](InferenceContext* c) {
      if (c->num_inputs() != 6) {
        return errors::InvalidArgument("len(inputs) != 6.");
      }
      // The following should stay in sync with `ComputeSparseToSparse` shape
      // assertions in kernels/set_kernels.cc.
      // Ranks must be compatible, and be >= 2.
      ShapeHandle input0_shape_shape = c->input(2);
      ShapeHandle input1_shape_shape = c->input(5);
      TF_RETURN_IF_ERROR(shape_inference::ValidateSparseTensor(
          c, c->input(0), c->input(1), input0_shape_shape));
      TF_RETURN_IF_ERROR(shape_inference::ValidateSparseTensor(
          c, c->input(3), c->input(4), input1_shape_shape));

      DimensionHandle input0_rank_dim = c->Dim(input0_shape_shape, 0);
      DimensionHandle input1_rank_dim = c->Dim(input1_shape_shape, 0);
      DimensionHandle output_rank_dim;
      if (c->ValueKnown(input0_rank_dim)) {
        const int64 input0_rank = c->Value(input0_rank_dim);
        if (input0_rank < 2) {
          return errors::InvalidArgument("Input 0, expected rank >= 2, got ",
                                         input0_rank, ".");
        }
        TF_RETURN_IF_ERROR(
            c->WithValue(input1_rank_dim, input0_rank, &input1_rank_dim));
        output_rank_dim = input0_rank_dim;
      } else if (c->ValueKnown(input1_rank_dim)) {
        const int64 input1_rank = c->Value(input1_rank_dim);
        if (input1_rank < 2) {
          return errors::InvalidArgument("Input 1, expected rank >= 2, got ",
                                         input1_rank, ".");
        }
        output_rank_dim = input1_rank_dim;
      } else {
        output_rank_dim = c->UnknownDim();
      }

      c->set_output(0, c->Matrix(c->UnknownDim(), output_rank_dim));
      c->set_output(1, c->Vector(c->UnknownDim()));
      c->set_output(2, c->Vector(output_rank_dim));
      return Status::OK();
    });

}  // namespace tensorflow