// Copyright 2019 The MediaPipe Authors.
//
// 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 "mediapipe/calculators/tensor/tensors_to_floats_calculator.pb.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/ret_check.h"

namespace mediapipe {

namespace {

inline float Sigmoid(float value) { return 1.0f / (1.0f + std::exp(-value)); }

}  // namespace

// A calculator for converting Tensors to to a float or a float vector.
//
// Input:
//  TENSORS - Vector of Tensors of type kFloat32. Only the first
//            tensor will be used.
// Output:
//  FLOAT(optional) - Converted single float number.
//  FLOATS(optional) - Converted float vector.
//
// Notes: To output FLOAT stream, the input tensor must have size 1, e.g.
//        only 1 float number in the tensor.
//
// Usage example:
// node {
//   calculator: "TensorsToFloatsCalculator"
//   input_stream: "TENSORS:tensors"
//   output_stream: "FLOATS:floats"
// }
namespace api2 {
class TensorsToFloatsCalculator : public Node {
 public:
  static constexpr Input<std::vector<Tensor>> kInTensors{"TENSORS"};
  static constexpr Output<float>::Optional kOutFloat{"FLOAT"};
  static constexpr Output<std::vector<float>>::Optional kOutFloats{"FLOATS"};
  MEDIAPIPE_NODE_INTERFACE(TensorsToFloatsCalculator, kInTensors, kOutFloat,
                           kOutFloats);

  static absl::Status UpdateContract(CalculatorContract* cc);
  absl::Status Open(CalculatorContext* cc) final;
  absl::Status Process(CalculatorContext* cc) final;

 private:
  ::mediapipe::TensorsToFloatsCalculatorOptions options_;
};
MEDIAPIPE_REGISTER_NODE(TensorsToFloatsCalculator);

absl::Status TensorsToFloatsCalculator::UpdateContract(CalculatorContract* cc) {
  // Only exactly a single output allowed.
  RET_CHECK(kOutFloat(cc).IsConnected() ^ kOutFloats(cc).IsConnected());
  return absl::OkStatus();
}

absl::Status TensorsToFloatsCalculator::Open(CalculatorContext* cc) {
  options_ = cc->Options<::mediapipe::TensorsToFloatsCalculatorOptions>();
  return absl::OkStatus();
}

absl::Status TensorsToFloatsCalculator::Process(CalculatorContext* cc) {
  const auto& input_tensors = *kInTensors(cc);
  RET_CHECK(!input_tensors.empty());
  RET_CHECK(input_tensors[0].element_type() == Tensor::ElementType::kFloat32);
  // TODO: Add option to specify which tensor to take from.
  auto view = input_tensors[0].GetCpuReadView();
  auto raw_floats = view.buffer<float>();
  int num_values = input_tensors[0].shape().num_elements();
  auto output_floats = absl::make_unique<std::vector<float>>(
      raw_floats, raw_floats + num_values);

  switch (options_.activation()) {
    case TensorsToFloatsCalculatorOptions::SIGMOID:
      std::transform(output_floats->begin(), output_floats->end(),
                     output_floats->begin(), Sigmoid);
      break;
    case TensorsToFloatsCalculatorOptions::NONE:
      break;
  }

  if (kOutFloat(cc).IsConnected()) {
    // TODO: Could add an index in the option to specifiy returning
    // one value of a float array.
    RET_CHECK_EQ(num_values, 1);
    kOutFloat(cc).Send(output_floats->at(0));
  } else {
    kOutFloats(cc).Send(std::move(output_floats));
  }
  return absl::OkStatus();
}
}  // namespace api2
}  // namespace mediapipe
