/*
 * Copyright (c) 2019 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 "TestHelper.h"

#include "Importer.h"

#include "IR/TFStopGradient.h"

#include <loco.h>
#include <plier/tf/TestHelper.h>

#include <gtest/gtest.h>

#include <tensorflow/core/framework/graph.pb.h>

using namespace moco::tf::test;

namespace
{
// clang-format off
const char *stopgradient_basic_pbtxt = STRING_CONTENT(
node {
  name: "Placeholder"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 1
        }
        dim {
          size: 2
        }
        dim {
          size: 3
        }
        dim {
          size: 4
        }
      }
    }
  }
}
node {
  name: "StopGradient_01"
  op: "StopGradient"
  input: "Placeholder"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
}
);
// clang-format on

} // namespace

TEST(TensorFlowImport, tf_stopgradient_basic)
{
  // load graph
  moco::tf::Importer importer;
  moco::tf::ModelSignature signature;
  signature.add_output(moco::tf::TensorName("StopGradient_01", 0));

  tensorflow::GraphDef graph_def;
  EXPECT_TRUE(plier::tf::parse_graphdef(stopgradient_basic_pbtxt, graph_def));
  std::unique_ptr<loco::Graph> graph = importer.import(signature, graph_def);

  // what to test:
  // - TFStopGradient node should exist
  // - input() should not be null

  auto node = moco::tf::test::find_first_node_bytype<moco::tf::TFStopGradient>(graph.get());

  ASSERT_NE(node, nullptr);
  ASSERT_NE(node->input(), nullptr);
}