/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include "gtest/gtest.h"
#include "ascir_ops.h"
#include "codegen_kernel.h"
#include "elewise/remove_pad_api_call.h"

using namespace af::ops;
using namespace codegen;
using namespace af::ascir_op;
using namespace testing;

class RemovePadApiCallTest : public ::testing::Test {
 protected:
  void SetUp() override {}
  void TearDown() override {}
};

TEST_F(RemovePadApiCallTest, RemovePadApiCall_01) {
  std::string binaryname = "removepad";

  std::vector<ascir::AxisId> current_axis;
  std::vector<std::reference_wrapper<const Tensor>> inputs;
  std::vector<std::reference_wrapper<const Tensor>> outputs;
  std::string result;

  af::SizeVar s0(af::Symbol("s0"));
  af::SizeVar s1(af::Symbol("s1"));
  af::SizeVar s2(af::Symbol("s2"));

  af::Axis z0{.id = 0, .name = "z0", .type = af::Axis::Type::kAxisTypeTileOuter, .size = s0.expr};
  af::Axis z1{.id = 1, .name = "z1", .type = af::Axis::Type::kAxisTypeTileInner, .size = s1.expr};
  af::Axis z2{.id = 2, .name = "z2", .type = af::Axis::Type::kAxisTypeOriginal, .size = s2.expr};

  codegen::Tiler tiler;
  tiler.AddSizeVar(af::SizeVar(s0));
  tiler.AddSizeVar(af::SizeVar(s1));
  tiler.AddSizeVar(af::SizeVar(s2));
  tiler.AddAxis(z0);
  tiler.AddAxis(z1);
  tiler.AddAxis(z2);

  codegen::TPipe tpipe("tpipe", tiler);
  af::AscGraph graph("test");
  af::ascir_op::Data x("x", graph);

  auto node = graph.FindNode("x");
  af::AscTensor tensor = node->outputs[0];

  tensor.attr.axis = {z0.id, z1.id, z2.id};
  tensor.attr.vectorized_axis = {z1.id, z2.id};
  tensor.attr.repeats = {z0.size, z1.size, z2.size};
  tensor.attr.strides = {z1.size * z2.size, z2.size, One};
  tensor.attr.mem.tensor_id = 1;
  tensor.attr.mem.alloc_type = af::AllocType::kAllocTypeBuffer;
  tensor.attr.mem.position = af::Position::kPositionVecIn;
  tensor.attr.opt.merge_scope = af::kIdNone;
  tensor.attr.buf.id = 2;
  vector<af::Expression> vectorized_strides{One, One};
  vectorized_strides[0] = z2.size;
  tensor.attr.vectorized_strides = vectorized_strides;

  std::string dtype_name;
  Tensor::DtypeName(tensor.attr.dtype, dtype_name);

  Tensor tensor1(tensor, dtype_name);
  tensor1.is_constant = true;
  tensor1.vectorized_axis_pos = {1, 2};
  tensor1.axis_size = {s0.expr, s1.expr, s2.expr};

  Tensor tensor2(tensor, dtype_name);
  tensor2.is_constant = true;
  tensor2.vectorized_axis_pos = {1, 2};
  tensor2.axis_size = {s0.expr, s1.expr, s2.expr};
  inputs.push_back(std::ref(tensor1));
  inputs.push_back(std::ref(tensor2));

  Tensor output_tensor(tensor, dtype_name);
  output_tensor.is_constant = true;
  output_tensor.vectorized_axis_pos = {1, 2};
  output_tensor.axis_size = {s0.expr, s1.expr, s2.expr};
  outputs.push_back(std::ref(output_tensor));
  codegen::RemovePadApiCall call(binaryname);
  Status status = call.Generate(tpipe, current_axis, inputs, outputs, result);
  ;
  EXPECT_EQ(status, af::SUCCESS);
}

TEST_F(RemovePadApiCallTest, RemovePadApiCall_02) {
  std::string binaryname = "removepad";

  std::vector<ascir::AxisId> current_axis;
  std::vector<std::reference_wrapper<const Tensor>> inputs;
  std::vector<std::reference_wrapper<const Tensor>> outputs;
  std::string result;

  af::SizeVar s0(af::Symbol("s0"));
  af::SizeVar s1(af::Symbol("s1"));

  af::Axis z0{.id = 0, .name = "z0", .type = af::Axis::Type::kAxisTypeTileOuter, .size = s0.expr};
  af::Axis z1{.id = 1, .name = "z1", .type = af::Axis::Type::kAxisTypeTileInner, .size = s1.expr};

  codegen::Tiler tiler;
  tiler.AddSizeVar(af::SizeVar(s0));
  tiler.AddSizeVar(af::SizeVar(s1));
  tiler.AddAxis(z0);
  tiler.AddAxis(z1);

  codegen::TPipe tpipe("tpipe", tiler);
  af::AscGraph graph("test");
  af::ascir_op::Data x("x", graph);

  auto node = graph.FindNode("x");
  af::AscTensor tensor = node->outputs[0];

  tensor.attr.axis = {z0.id, z1.id};
  tensor.attr.vectorized_axis = {z1.id};
  tensor.attr.repeats = {z0.size, z1.size};
  tensor.attr.strides = {z1.size, One};
  tensor.attr.mem.tensor_id = 1;
  tensor.attr.mem.alloc_type = af::AllocType::kAllocTypeBuffer;
  tensor.attr.mem.position = af::Position::kPositionVecIn;
  tensor.attr.opt.merge_scope = af::kIdNone;
  tensor.attr.buf.id = 2;
  vector<af::Expression> vectorized_strides{One, One};
  vectorized_strides[0] = z1.size;
  tensor.attr.vectorized_strides = vectorized_strides;

  std::string dtype_name;
  Tensor::DtypeName(tensor.attr.dtype, dtype_name);

  Tensor tensor1(tensor, dtype_name);
  tensor1.is_constant = true;
  tensor1.vectorized_axis_pos = {1};
  tensor1.axis_size = {s0.expr, s1.expr};

  Tensor tensor2(tensor, dtype_name);
  tensor2.is_constant = true;
  tensor2.vectorized_axis_pos = {1};
  tensor2.axis_size = {s0.expr, s1.expr};
  inputs.push_back(std::ref(tensor1));
  inputs.push_back(std::ref(tensor2));

  Tensor output_tensor(tensor, dtype_name);
  output_tensor.is_constant = true;
  output_tensor.vectorized_axis_pos = {1};
  output_tensor.axis_size = {s0.expr, s1.expr};
  outputs.push_back(std::ref(output_tensor));

  codegen::RemovePadApiCall call(binaryname);
  Status status = call.Generate(tpipe, current_axis, inputs, outputs, result);
  ;
  EXPECT_EQ(status, af::SUCCESS);
}