// Copyright (c) ONNX Project Contributors
//
// SPDX-License-Identifier: Apache-2.0

#include "onnx_core/backend_test/expect.h"
#include "onnx_core/runtime/kernels/random.h"
#include "onnx_extensions/backend_test/cases/tensor/include_tensor_cases.h"
#include "onnx_extensions/kernels/kernels/tensor/include_tensor_kernels.h"
#include "onnx_proto/onnx_helper.h"

#include <cstdint>
#include <string>
#include <vector>

namespace ONNX_LIGHT_NAMESPACE::onnx_backend_test {

namespace {

NodeProto MakeTransposeNode(const std::vector<int64_t> &perm = {}) {
  NodeProto node;
  node.set_op_type("Transpose");
  node.add_input("data");
  node.add_output("transposed");
  if (!perm.empty()) {
    AddAttribute<std::vector<int64_t>>(node, "perm", perm);
  }
  return node;
}

} // namespace

void RegisterTransposeCases(std::vector<TestCase> &registry, TestMode mode) {
  const OpsetId opset = DefaultOpset(13);
  const KernelContext ctx{opset};
  const onnx_kernels::kernel::Transpose transpose_kernel{ctx};

  if (mode == TestMode::BENCHMARK) {
    NodeProto node = MakeTransposeNode();
    Expect(registry, std::move(node), "test_cc_transpose_default_perm_benchmark", {opset},
           {4194304}, {4194304}, [transpose_kernel]() -> IoData {
             Tensor data = RandnTensor(DataType::FLOAT, {2048, 2048}, 2001);
             Tensor transposed = transpose_kernel(data, /*perm=*/{});
             return IoData{{std::move(data)}, {std::move(transposed)}};
           });
    return;
  }

  // test_cc_transpose_default_perm
  {
    Expect(registry, MakeTransposeNode(), "test_cc_transpose_default_perm", {opset},
           [=]() -> IoData {
             const Tensor data = Tensor::FromFloat("", {2, 3}, {0.f, 1.f, 2.f, 3.f, 4.f, 5.f});
             const Tensor transposed = transpose_kernel(data, /*perm=*/{});
             return IoData{{std::move(data)}, {std::move(transposed)}};
           });
  }

  // test_cc_transpose_permuted_axes
  {
    const std::vector<int64_t> perm{1, 0, 2};
    Expect(registry, MakeTransposeNode(perm), "test_cc_transpose_permuted_axes", {opset},
           [=]() -> IoData {
             const Tensor data =
                 Tensor::FromFloat("", {2, 3, 4}, {0.f,  1.f,  2.f,  3.f,  4.f,  5.f,  6.f,  7.f,
                                                   8.f,  9.f,  10.f, 11.f, 12.f, 13.f, 14.f, 15.f,
                                                   16.f, 17.f, 18.f, 19.f, 20.f, 21.f, 22.f, 23.f});
             const Tensor transposed = transpose_kernel(data, perm);
             return IoData{{std::move(data)}, {std::move(transposed)}};
           });
  }

  // test_cc_transpose_permuted_axes_2
  {
    const std::vector<int64_t> perm{1, 2, 0};
    Expect(registry, MakeTransposeNode(perm), "test_cc_transpose_permuted_axes_2", {opset},
           [=]() -> IoData {
             const Tensor data = Tensor::FromFloat("", {1, 2, 3}, {0.f, 1.f, 2.f, 3.f, 4.f, 5.f});
             const Tensor transposed = transpose_kernel(data, perm);
             return IoData{{std::move(data)}, {std::move(transposed)}};
           });
  }

  // Mirror upstream ``Transpose.export_all_permutations``: iterate over every
  // permutation of the 3-axis input shape ``(2, 3, 4)``. The upstream test
  // suite registers six cases ``transpose_all_permutations_{0..5}``. Each
  // light case must contain that suffix as a substring of its name so the
  // backend-test name comparison treats them as covered.
  {
    const std::vector<int64_t> shape{2, 3, 4};
    const std::vector<std::vector<int64_t>> all_perms = {
        {0, 1, 2}, {0, 2, 1}, {1, 0, 2}, {1, 2, 0}, {2, 0, 1}, {2, 1, 0},
    };
    for (size_t i = 0; i < all_perms.size(); ++i) {
      const std::vector<int64_t> perm = all_perms[i];
      const std::string name = "test_cc_transpose_all_permutations_" + std::to_string(i);
      Expect(registry, MakeTransposeNode(perm), name, {opset},
             [transpose_kernel, shape, perm]() -> IoData {
               Tensor data = RandnTensor(DataType::FLOAT, shape, /*seed=*/41);
               Tensor transposed = transpose_kernel(data, perm);
               return IoData{{std::move(data)}, {std::move(transposed)}};
             });
    }
  }
}

} // namespace ONNX_LIGHT_NAMESPACE::onnx_backend_test
