/*
 * 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 "BiasAdd.h"
#include "TestHelper.h"

#include "Importer.h"
#include "IR/TFBiasAdd.h"

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

#include <gtest/gtest.h>

using namespace moco::tf;
using namespace moco::tf::test;

namespace
{

// clang-format off
const char *bias_add_01_pbtxtdata = STRING_CONTENT(

node {
  name: "val"
  op: "Const"
  attr {
    key: "dtype"
    value { type: DT_FLOAT }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_FLOAT
        tensor_shape {
          dim { size: 1 }
          dim { size: 5 }
          dim { size: 5 }
          dim { size: 3 }
        }
        float_val: 2.1
      }
    }
  }
}
node {
  name: "bias"
  op: "Const"
  attr {
    key: "dtype"
    value { type: DT_FLOAT }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_FLOAT
        tensor_shape {
          dim { size: 3 }
        }
        float_val: 1.1
      }
    }
  }
}
node {
  name: "out"
  op: "BiasAdd"
  input: "val"
  input: "bias"
  attr {
    key: "T"
    value { type: DT_FLOAT }
  }
  attr {
    key: "data_format"
    value { s: "NHWC" }
  }
}

);
// clang-format on

} // namespace

TEST(TensorFlowImport, bias_add_01)
{
  moco::tf::ModelSignature signature;
  signature.add_output(moco::tf::TensorName("out", 0));

  tensorflow::GraphDef graph_def;
  EXPECT_TRUE(plier::tf::parse_graphdef(bias_add_01_pbtxtdata, graph_def));

  // Test "BiasAddGraphBuilderImpl<ImportTarget::TensorFlow>"
  {
    using BiasAddGraphBuilder = BiasAddGraphBuilderImpl<ImportTarget::TensorFlow>;

    moco::tf::GraphBuilderRegistry r{&moco::tf::GraphBuilderRegistry::get()};
    r.add("BiasAdd", stdex::make_unique<BiasAddGraphBuilder>());
    moco::tf::Importer importer{&r};

    std::unique_ptr<loco::Graph> graph = importer.import(signature, graph_def);

    // what to test:
    // - there should exist TFBiasAdd
    // - value() should not be nullptr
    // - bias() should not be nullptr
    // - data_layout should match

    moco::tf::TFBiasAdd *bias_add =
        moco::tf::test::find_first_node_bytype<moco::tf::TFBiasAdd>(graph.get());

    ASSERT_NE(bias_add, nullptr);

    ASSERT_NE(bias_add->value(), nullptr);
    ASSERT_NE(bias_add->bias(), nullptr);

    ASSERT_TRUE(bias_add->data_layout() == "NHWC");
  }
  // Test "BiasAddGraphBuilderImpl<ImportTarget::Canonical>"
  {
    using BiasAddGraphBuilder = BiasAddGraphBuilderImpl<ImportTarget::Canonical>;

    moco::tf::GraphBuilderRegistry r{&moco::tf::GraphBuilderRegistry::get()};
    r.add("BiasAdd", stdex::make_unique<BiasAddGraphBuilder>());
    moco::tf::Importer importer{&r};

    std::unique_ptr<loco::Graph> graph = importer.import(signature, graph_def);

    // what to test:
    // - there should exist BiasAdd
    // - value() should not be nullptr
    // - bias() input should be BiasEncode
    // - axis should match

    // loco node : ------------+-- BiasAdd<Domain::Tensor> --
    //             BiasEncode -/

    loco::BiasAdd<loco::Domain::Tensor> *bias_add =
        moco::tf::test::find_first_node_bytype<loco::BiasAdd<loco::Domain::Tensor>>(graph.get());

    ASSERT_NE(bias_add, nullptr);
    ASSERT_NE(bias_add->value(), nullptr);

    auto bias_enc = dynamic_cast<loco::BiasEncode *>(bias_add->bias());
    ASSERT_NE(bias_enc, nullptr);

    ASSERT_EQ(bias_add->axis(), 3); // NHWC
  }
}

namespace
{

// clang-format off
const char *bias_add_NCHW_pbtxtdata = STRING_CONTENT(

node {
  name: "val"
  op: "Const"
  attr {
    key: "dtype"
    value { type: DT_FLOAT }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_FLOAT
        tensor_shape {
          dim { size: 1 }
          dim { size: 3 }
          dim { size: 299 }
          dim { size: 299 }
        }
        float_val: 2.1
      }
    }
  }
}
node {
  name: "bias"
  op: "Const"
  attr {
    key: "dtype"
    value { type: DT_FLOAT }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_FLOAT
        tensor_shape {
          dim { size: 3 }
        }
        float_val: 1.1
      }
    }
  }
}
node {
  name: "out"
  op: "BiasAdd"
  input: "val"
  input: "bias"
  attr {
    key: "T"
    value { type: DT_FLOAT }
  }
  attr {
    key: "data_format"
    value { s: "NCHW" }
  }
}

);
// clang-format on

} // namespace

TEST(TensorFlowImport, bias_add_NCHW_axis)
{
  moco::tf::Importer importer;
  moco::tf::ModelSignature signature;
  signature.add_output(moco::tf::TensorName("out", 0));

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

  // Test "BiasAddGraphBuilderImpl<ImportTarget::TensorFlow>"
  {
    using BiasAddGraphBuilder = BiasAddGraphBuilderImpl<ImportTarget::TensorFlow>;

    moco::tf::GraphBuilderRegistry r{&moco::tf::GraphBuilderRegistry::get()};
    r.add("BiasAdd", stdex::make_unique<BiasAddGraphBuilder>());
    moco::tf::Importer importer{&r};

    std::unique_ptr<loco::Graph> graph = importer.import(signature, graph_def);

    // what to test:
    // - there should exist TFBiasAdd
    // - value() should not be nullptr
    // - bias() should not be nullptr
    // - data_layout should match

    moco::tf::TFBiasAdd *bias_add =
        moco::tf::test::find_first_node_bytype<moco::tf::TFBiasAdd>(graph.get());

    ASSERT_NE(bias_add, nullptr);

    ASSERT_NE(bias_add->value(), nullptr);
    ASSERT_NE(bias_add->bias(), nullptr);

    ASSERT_TRUE(bias_add->data_layout() == "NCHW");
  }
  // Test "BiasAddGraphBuilderImpl<ImportTarget::Canonical>"
  {
    using BiasAddGraphBuilder = BiasAddGraphBuilderImpl<ImportTarget::Canonical>;

    moco::tf::GraphBuilderRegistry r{&moco::tf::GraphBuilderRegistry::get()};
    r.add("BiasAdd", stdex::make_unique<BiasAddGraphBuilder>());
    moco::tf::Importer importer{&r};

    std::unique_ptr<loco::Graph> graph = importer.import(signature, graph_def);

    // what to test:
    // - there should exist BiasAdd
    // - value() should not be nullptr
    // - bias() input should be BiasEncode
    // - axis should match

    // loco node : ------------+-- BiasAdd<Domain::Tensor> --
    //             BiasEncode -/

    loco::BiasAdd<loco::Domain::Tensor> *bias_add =
        moco::tf::test::find_first_node_bytype<loco::BiasAdd<loco::Domain::Tensor>>(graph.get());

    ASSERT_NE(bias_add, nullptr);
    ASSERT_NE(bias_add->value(), nullptr);

    auto bias_enc = dynamic_cast<loco::BiasEncode *>(bias_add->bias());
    ASSERT_NE(bias_enc, nullptr);

    ASSERT_EQ(bias_add->axis(), 1); // NCHW
  }
}