/*
 * Copyright (c) 2018 Samsung Electronics Co., Ltd. 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 <gtest/gtest.h>

#include "graph/Graph.h"
#include "model/Model.h"
#include "model/Index.h"
#include "model/OperandIndexSequence.h"
#include "model/operation/Conv2DNode.h"
#include "model/operation/ConcatNode.h"

#include <cpp14/memory.h>

#include <stdexcept>

using Index = neurun::model::IOIndex;
using IndexSet = neurun::model::OperandIndexSequence;

TEST(graph_operation_setIO, operation_setIO_conv)
{
  neurun::model::Model model;

  neurun::model::Shape shape{3};
  neurun::model::TypeInfo type{neurun::model::DataType::INT32};

  // Add Conv
  using GraphNode = neurun::model::operation::Conv2DNode;

  auto input_operand = model.operands.emplace(shape, type);
  auto kernel_operand = model.operands.emplace(shape, type);
  auto bias_operand = model.operands.emplace(shape, type);
  IndexSet inputs{input_operand, kernel_operand, bias_operand};

  GraphNode::Param conv_params;
  conv_params.padding.type = neurun::model::PaddingType::SAME;
  conv_params.stride.horizontal = 1;
  conv_params.stride.vertical = 1;
  conv_params.activation = neurun::model::Activation::NONE;

  auto output_operand = model.operands.emplace(shape, type).value();
  IndexSet outputs{output_operand};

  auto conv = nnfw::cpp14::make_unique<GraphNode>(inputs, outputs, conv_params);

  ASSERT_NE(conv, nullptr);
  ASSERT_EQ(conv->getInputs().at(Index{0}).value(), inputs.at(0).value());
  conv->setInputs({8, 9, 10});
  ASSERT_NE(conv->getInputs().at(Index{0}).value(), inputs.at(0).value());
  ASSERT_EQ(conv->getInputs().at(Index{0}).value(), 8);
}

TEST(graph_operation_setIO, operation_setIO_concat)
{
  neurun::model::Model model;

  neurun::model::Shape shape{3};

  neurun::model::TypeInfo type{neurun::model::DataType::INT32};

  using GraphNode = neurun::model::operation::ConcatNode;

  // Add Concat
  IndexSet inputs;
  for (int i = 0; i < 6; ++i)
  {
    inputs.append(model.operands.emplace(shape, type));
  }

  GraphNode::Param concat_params{0};

  auto output_operand = model.operands.emplace(shape, type).value();
  IndexSet outputs{output_operand};

  auto concat = nnfw::cpp14::make_unique<GraphNode>(inputs, outputs, concat_params);

  ASSERT_NE(concat, nullptr);
  ASSERT_EQ(concat->getInputs().size(), 6);
  ASSERT_EQ(concat->getInputs().at(Index{0}).value(), inputs.at(0).value());

  concat->setInputs({80, 6, 9, 11});
  ASSERT_EQ(concat->getInputs().size(), 4);
  ASSERT_NE(concat->getInputs().at(Index{0}).value(), inputs.at(0).value());
  ASSERT_EQ(concat->getInputs().at(Index{0}).value(), 80);
  ASSERT_EQ(concat->getInputs().at(Index{2}).value(), 9);
  ASSERT_THROW(concat->getInputs().at(Index{5}), std::out_of_range);
}