/* Copyright 2015 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.
==============================================================================*/

// See docs in ../ops/data_flow_ops.cc.

#include <deque>
#include <vector>

#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/kernels/padding_fifo_queue.h"
#include "tensorflow/core/kernels/queue_base.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/batch_util.h"

namespace tensorflow {

PaddingFIFOQueue::PaddingFIFOQueue(
    int capacity, const DataTypeVector& component_dtypes,
    const std::vector<PartialTensorShape>& partial_shapes, const string& name)
    : FIFOQueue(capacity, component_dtypes,
                ConvertShapesPartialDimensionsToZero(partial_shapes), name),
      partial_shapes_(partial_shapes) {}

Status PaddingFIFOQueue::Initialize() {
  Status s = FIFOQueue::Initialize();
  if (!s.ok()) return s;

  if (component_dtypes_.size() != partial_shapes_.size()) {
    return errors::InvalidArgument(
        "Shapes must be provided for all components, but received ",
        component_dtypes_.size(), " dtypes and ", partial_shapes_.size(),
        " shapes.");
  }

  return Status::OK();
}

/* static */
Status PaddingFIFOQueue::GetElementComponent(
    const PaddingFIFOQueue::Tuple& tuple, int component, OpKernelContext* ctx,
    PersistentTensor* out_tensor) {
  TensorShape element_shape(tuple[component].shape());
  Tensor* element_access = nullptr;
  TF_RETURN_IF_ERROR(ctx->allocate_persistent(
      tuple[component].dtype(), element_shape, out_tensor, &element_access));
  *element_access = tuple[component];
  return Status::OK();
}

void PaddingFIFOQueue::TryDequeueMany(int num_elements, OpKernelContext* ctx,
                                      bool allow_small_batch,
                                      CallbackWithTuple callback) {
  if (num_elements == 0) {
    Tuple tuple;
    tuple.reserve(num_components());
    for (int i = 0; i < num_components(); ++i) {
      // TODO(josh11b,misard): Switch to allocate_output().
      // See similar comment in fifo_queue.cc
      Tensor element;
      // Here, ManyOutShape returns zeros for undetermined shapes,
      // which is exactly what we want to use.
      OP_REQUIRES_OK(ctx, ctx->allocate_temp(component_dtypes_[i],
                                             ManyOutShape(i, 0), &element));
      tuple.emplace_back(element);
    }
    callback(tuple);
    return;
  }

  CancellationManager* cm = ctx->cancellation_manager();
  CancellationToken token = cm->get_cancellation_token();
  bool already_cancelled;
  {
    mutex_lock l(mu_);
    already_cancelled = !cm->RegisterCallback(
        token, [this, cm, token]() { Cancel(kDequeue, cm, token); });
    if (!already_cancelled) {
      // TODO(josh11b): This makes two copies of callback, avoid this if possible.
      dequeue_attempts_.emplace_back(
          num_elements, [callback]() { callback(Tuple()); }, ctx, cm, token,
          [callback, allow_small_batch,
           this](Attempt* attempt) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
            int32 queue_size = queues_[0].size();
            if (closed_ && queue_size < attempt->elements_requested) {
              // If we don't have enough for a full dequeue, we have
              // to reset the attempt tuple.
              if (!attempt->tuples.empty()) {
                // Restore already-dequeued elements to the front of the queue.
                for (int64 i = attempt->tuples.size() - 1; i >= 0; --i) {
                  for (int j = 0; j < num_components(); ++j) {
                    PersistentTensor element;
                    Status s = GetElementComponent(attempt->tuples[i], j,
                                                   attempt->context, &element);
                    if (!s.ok()) {
                      attempt->context->SetStatus(
                          errors::DataLoss("Failed to restore element from "
                                           "partially-dequeued batch "
                                           "to PaddingFIFOQueue: ",
                                           s.error_message()));
                    }
                    queues_[j].push_front(element);
                  }
                }
              }
              if (allow_small_batch && !queues_[0].empty()) {
                // Request all remaining elements in the queue.
                queue_size = queues_[0].size();
                attempt->tuples.clear();
                attempt->elements_requested = queue_size;
              } else {
                if (allow_small_batch) {
                  // There may be some enqueue attempts containing
                  // values.  If so, we'll yield and wait for them
                  // to add elements to the queue.
                  if (!enqueue_attempts_.empty()) return kProgress;
                }
                if (attempt->context->status().ok()) {
                  attempt->context->SetStatus(errors::OutOfRange(
                      "PaddingFIFOQueue '", name_, "' is closed and has ",
                      "insufficient elements (requested ",
                      attempt->elements_requested, ", current size ",
                      queue_size, ")"));
                }
                return kComplete;
              }
            }

            RunResult result = kNoProgress;
            for (; queue_size > 0; --queue_size) {
              result = kProgress;
              Tuple tuple;
              DequeueLocked(attempt->context, &tuple);
              attempt->tuples.push_back(tuple);
              tuple.clear();
              --attempt->elements_requested;

              if (attempt->elements_requested == 0) {
                // Finished.  Allocate attempt->tuple and
                // copy from attempt->tuples to attempt->tuple.
                attempt->tuple.reserve(num_components());
                std::vector<Tuple>& tuples = attempt->tuples;

                std::vector<bool> dynamic_shape;
                const int64 batch_size = tuples.size();

                for (int i = 0; i < num_components(); ++i) {
                  const PartialTensorShape partial_shape =
                      PartialTensorShape({batch_size})
                          .Concatenate(partial_shapes_[i]);
                  TensorShape shape({batch_size});

                  for (int j = 0; j < partial_shape.dims() - 1; ++j) {
                    if (partial_shape.dim_size(j + 1) > -1) {
                      shape.AddDim(partial_shape.dim_size(j + 1));
                    } else {
                      // Expand sizes to match.
                      int64 max_val = 0;
                      for (const Tuple& t : tuples) {
                        max_val = std::max(max_val, t[i].shape().dim_size(j));
                      }
                      shape.AddDim(max_val);
                    }
                  }

                  Tensor element;
                  attempt->context->SetStatus(attempt->context->allocate_temp(
                      component_dtypes_[i], shape, &element));
                  if (!attempt->context->status().ok()) return kComplete;

                  bool has_dynamic_shape = !partial_shape.IsFullyDefined();
                  if (has_dynamic_shape) {
                    // Set all values to zero because not all values
                    // will get written over.
                    attempt->context->SetStatus(SetElementZero(&element));
                    if (!attempt->context->status().ok()) return kComplete;
                  }

                  dynamic_shape.push_back(has_dynamic_shape);

                  // TODO(ebrevdo): should this be a persistent tensor?
                  attempt->tuple.emplace_back(element);
                }

                for (size_t index = 0; index < tuples.size(); ++index) {
                  for (int i = 0; i < num_components(); ++i) {
                    if (dynamic_shape[i]) {
                      // Slightly slower copy operation
                      attempt->context->SetStatus(CopyElementToLargerSlice(
                          tuples[index][i], &attempt->tuple[i], index));
                    } else {
                      attempt->context->SetStatus(
                          batch_util::CopyElementToSlice(
                              std::move(tuples[index][i]), &attempt->tuple[i],
                              index));
                    }
                    if (!attempt->context->status().ok()) return kComplete;
                  }
                }
                tuple = attempt->tuple;
                attempt->tuples.clear();
                attempt->done_callback = [callback, tuple]() {
                  callback(tuple);
                };
                return kComplete;
              }
            }
            return result;
          });
    }
  }
  if (!already_cancelled) {
    FlushUnlocked();
  } else {
    ctx->SetStatus(errors::Cancelled("Dequeue operation was cancelled"));
    callback(Tuple());
  }
}

Status PaddingFIFOQueue::ValidateTuple(const Tuple& tuple) {
  TF_RETURN_IF_ERROR(ValidateTupleCommon(tuple));
  for (size_t i = 0; i < tuple.size(); ++i) {
    if (!partial_shapes_[i].IsCompatibleWith(tuple[i].shape())) {
      return errors::InvalidArgument("Shape mismatch in tuple component ", i,
                                     ". Expected ",
                                     partial_shapes_[i].DebugString(), ", got ",
                                     tuple[i].shape().DebugString());
    }
  }
  return Status::OK();
}

Status PaddingFIFOQueue::ValidateManyTuple(const Tuple& tuple) {
  TF_RETURN_IF_ERROR(ValidateTupleCommon(tuple));
  const int64 batch_size = tuple[0].dim_size(0);
  for (size_t i = 0; i < tuple.size(); ++i) {
    // Expected shape is [batch_size] + partial_shapes_[i]
    const PartialTensorShape expected_shape =
        PartialTensorShape({batch_size}).Concatenate(partial_shapes_[i]);
    if (!expected_shape.IsCompatibleWith(tuple[i].shape())) {
      return errors::InvalidArgument("Shape mismatch in tuple component ", i,
                                     ". Expected ",
                                     expected_shape.DebugString(), ", got ",
                                     tuple[i].shape().DebugString());
    }
  }
  return Status::OK();
}

Status PaddingFIFOQueue::CompatibleNodeDefShapes(
    const NodeDef& node_def) const {
  std::vector<PartialTensorShape> requested_shapes;
  TF_RETURN_IF_ERROR(GetNodeAttr(node_def, "shapes", &requested_shapes));
  if (!PartialTensorShapeUtils::AreCompatible(requested_shapes,
                                              partial_shapes_)) {
    return errors::InvalidArgument(
        "Shared queue '", name_, "' has component shapes ",
        PartialTensorShapeUtils::PartialShapeListString(partial_shapes_),
        " but requested component shapes were ",
        PartialTensorShapeUtils::PartialShapeListString(requested_shapes));
  } else {
    return Status::OK();
  }
}

Status PaddingFIFOQueue::MatchesNodeDef(const NodeDef& node_def) {
  if (!MatchesNodeDefOp(node_def, "PaddingFIFOQueue").ok() &&
      !MatchesNodeDefOp(node_def, "PaddingFIFOQueueV2").ok()) {
    return errors::InvalidArgument("Expected PaddingFIFOQueue, found ",
                                   node_def.op());
  }
  TF_RETURN_IF_ERROR(MatchesNodeDefCapacity(node_def, capacity_));
  TF_RETURN_IF_ERROR(MatchesNodeDefTypes(node_def));
  TF_RETURN_IF_ERROR(CompatibleNodeDefShapes(node_def));
  return Status::OK();
}

static Status ValidateElementToLargerSlice(const Tensor& element,
                                           Tensor* parent) {
  DCHECK_NE(parent->dim_size(0), 0);
  if (element.NumElements() > (parent->NumElements() / parent->dim_size(0))) {
    TensorShape chip_shape = parent->shape();
    chip_shape.RemoveDim(0);
    return errors::Internal(
        "HandleElementToLargerSlice Cannot copy slice: number of entries in "
        "element is greater than number of elements in parent slice.  ",
        "Shapes are: [element]: ", element.shape().DebugString(),
        ", [parent slice]: ", chip_shape.DebugString());
  }
  return Status::OK();
}

template <typename T, int NDIMS>
Status HandleElementToLargerSlice(const Tensor& element, Tensor* parent,
                                  int index) {
  Status s = ValidateElementToLargerSlice(element, parent);
  if (!s.ok()) {
    return s;
  }
  if (element.NumElements() == 0) {
    return Status::OK();
  }
  auto element_t = element.tensor<T, NDIMS>();
  auto parent_t = parent->tensor<T, NDIMS + 1>();
  Eigen::DSizes<Eigen::DenseIndex, NDIMS + 1> slice_indices;
  slice_indices[0] = index;
  Eigen::DSizes<Eigen::DenseIndex, NDIMS + 1> slice_size;
  slice_size[0] = 1;
  for (size_t i = 1; i < slice_size.size(); ++i) {
    slice_size[i] = element_t.dimension(i - 1);
  }
  parent_t.slice(slice_indices, slice_size) = element_t.reshape(slice_size);
  return Status::OK();
}

namespace {

template <int NDIMS>
Status HandleElementToLargerSliceWithRank(const Tensor& element, Tensor* parent,
                                          int index) {
#define HANDLE_TYPE(T)                                                   \
  case DataTypeToEnum<T>::value: {                                       \
    return HandleElementToLargerSlice<T, NDIMS>(element, parent, index); \
  }

  switch (element.dtype()) {
    TF_CALL_ALL_TYPES(HANDLE_TYPE);
#undef HANDLE_TYPE
    default:
      return errors::Unimplemented(
          "HandleElementToLargerSliceWithRank Unhandled data type: ",
          DataTypeString(element.dtype()));
  }
}

}  // namespace

Status PaddingFIFOQueue::CopyElementToLargerSlice(const Tensor& element,
                                                  Tensor* parent, int index) {
  if (parent->dims() != element.dims() + 1) {
    return errors::Internal(
        "Mismatched ranks.  Element's rank is: ", element.dims(),
        " but element is meant to be a slice in output Tensor having rank: ",
        parent->dims(), " (should be: ", element.dims() + 1, ")");
  }

#define HANDLE_DIMS(NDIMS)                                                  \
  case NDIMS: {                                                             \
    TF_RETURN_IF_ERROR(                                                     \
        HandleElementToLargerSliceWithRank<NDIMS>(element, parent, index)); \
    return Status::OK();                                                    \
  }

  switch (element.dims()) {
    HANDLE_DIMS(0);
    HANDLE_DIMS(1);
    HANDLE_DIMS(2);
    HANDLE_DIMS(3);
    HANDLE_DIMS(4);
#undef HANDLE_DIMS
    default:
      return errors::Unimplemented("CopyElementToLargerSlice Unhandled rank: ",
                                   element.dims());
  }
}

// Static method
Status PaddingFIFOQueue::SetElementZero(Tensor* element) {
#define HANDLE_TYPE(T)                                \
  if (element->dtype() == DataTypeToEnum<T>::value) { \
    element->flat<T>().setConstant(T());              \
    return Status::OK();                              \
  }
  TF_CALL_ALL_TYPES(HANDLE_TYPE);
#undef HANDLE_TYPE
  return errors::Unimplemented("SetElementZero Unhandled data type: ",
                               DataTypeString(element->dtype()));
}

std::vector<TensorShape> PaddingFIFOQueue::ConvertShapesPartialDimensionsToZero(
    const gtl::ArraySlice<PartialTensorShape>& partial_shapes) {
  std::vector<TensorShape> shapes(partial_shapes.size());
  for (size_t i = 0; i < shapes.size(); ++i) {
    const PartialTensorShape& partial = partial_shapes[i];
    TensorShape& shape = shapes[i];
    for (int64 s : partial.dim_sizes()) shape.AddDim(s < 0 ? 0 : s);
  }
  return shapes;
}

}  // namespace tensorflow