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

#include "onnx_core/backend_test/expect.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 <cstring>
#include <string>
#include <vector>

namespace ONNX_LIGHT_NAMESPACE::onnx_backend_test {

// ---------------------------------------------------------------------------
// Upsample — nearest/linear interpolation, output_dim[i] = floor(input_dim[i]
// * scales[i]). Available in opsets 1, 7, 9 and 10 (deprecated since opset 10
// in favour of ``Resize``). Mirrors the upstream upsample node test in
// ``onnx/backend/test/case/node/upsample.py``.
// ---------------------------------------------------------------------------

namespace {

NodeProto MakeUpsampleNode(const std::string &mode) {
  NodeProto node;
  node.set_op_type("Upsample");
  node.add_input("X");
  node.add_input("scales");
  node.add_output("Y");
  if (!mode.empty()) {
    AddAttribute<std::string>(node, "mode", mode);
  }
  return node;
}

// Builds a 1-D FLOAT ``scales`` tensor from the given values.
Tensor MakeScalesTensor(const std::vector<float> &scales) {
  return Tensor::FromFloat("", {static_cast<int64_t>(scales.size())}, scales);
}

} // namespace

void RegisterUpsampleCases(std::vector<TestCase> &registry, TestMode mode) {
  const OpsetId opset = DefaultOpset(9);
  const KernelContext ctx{opset};
  const onnx_kernels::kernel::Upsample upsample_kernel{ctx};

  if (mode == TestMode::BENCHMARK) {
    NodeProto node = MakeUpsampleNode("nearest");
    Expect(registry, std::move(node), "test_cc_upsample_nearest_benchmark", {opset}, {1048576, 4},
           {6291456}, [upsample_kernel]() -> IoData {
             Tensor X = RandnTensor(DataType::FLOAT, {1, 1, 1024, 1024}, 2001);
             Tensor scales = MakeScalesTensor({1.0f, 1.0f, 2.0f, 3.0f});
             onnx_kernels::kernel::Upsample::Attributes attrs;
             attrs.mode = "nearest";
             Tensor Y = upsample_kernel(X, scales, attrs);
             return IoData{{std::move(X), std::move(scales)}, {std::move(Y)}};
           });
    return;
  }

  // test_cc_upsample_nearest — mirrors the upstream ``test_upsample_nearest``
  // node test exactly. NCHW input of shape [1, 1, 2, 2], scales
  // [1, 1, 2, 3], nearest mode -> output shape [1, 1, 4, 6].
  {
    Expect(registry, MakeUpsampleNode("nearest"), "test_cc_upsample_nearest", {opset},
           [=]() -> IoData {
             const Tensor X = Tensor::FromFloat("", {1, 1, 2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
             const Tensor scales = MakeScalesTensor({1.0f, 1.0f, 2.0f, 3.0f});
             onnx_kernels::kernel::Upsample::Attributes attrs;
             attrs.mode = "nearest";
             const Tensor Y = upsample_kernel(X, scales, attrs);
             return IoData{{std::move(X), std::move(scales)}, {std::move(Y)}};
           });
  }

  // test_cc_upsample_nearest_default_mode — mode attribute omitted (defaults
  // to "nearest" per the upstream spec).
  {
    Expect(registry, MakeUpsampleNode(""), "test_cc_upsample_nearest_default_mode", {opset},
           [=]() -> IoData {
             const Tensor X = Tensor::FromFloat("", {1, 1, 2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
             const Tensor scales = MakeScalesTensor({1.0f, 1.0f, 2.0f, 2.0f});
             onnx_kernels::kernel::Upsample::Attributes attrs;
             attrs.mode = "nearest";
             const Tensor Y = upsample_kernel(X, scales, attrs);
             return IoData{{std::move(X), std::move(scales)}, {std::move(Y)}};
           });
  }

  // test_cc_upsample_nearest_1d — minimal 1-D nearest-neighbor case, useful
  // as an independent sanity check for the strided index mapping.
  {
    Expect(registry, MakeUpsampleNode("nearest"), "test_cc_upsample_nearest_1d", {opset},
           [=]() -> IoData {
             const Tensor X = Tensor::FromFloat("", {3}, {10.0f, 20.0f, 30.0f});
             const Tensor scales = MakeScalesTensor({2.0f});
             onnx_kernels::kernel::Upsample::Attributes attrs;
             attrs.mode = "nearest";
             const Tensor Y = upsample_kernel(X, scales, attrs);
             return IoData{{std::move(X), std::move(scales)}, {std::move(Y)}};
           });
  }

  // test_cc_upsample_linear — 4-D bilinear case (NCHW, scales == 1 on N and
  // C). Output is computed by the same kernel and embedded as the reference
  // so the test exercises kernel determinism rather than numerical parity
  // with another implementation.
  {
    Expect(registry, MakeUpsampleNode("linear"), "test_cc_upsample_linear", {opset},
           [=]() -> IoData {
             const Tensor X = Tensor::FromFloat("", {1, 1, 2, 2}, {1.0f, 2.0f, 3.0f, 4.0f});
             const Tensor scales = MakeScalesTensor({1.0f, 1.0f, 2.0f, 2.0f});
             onnx_kernels::kernel::Upsample::Attributes attrs;
             attrs.mode = "linear";
             const Tensor Y = upsample_kernel(X, scales, attrs);
             return IoData{{std::move(X), std::move(scales)}, {std::move(Y)}};
           });
  }
}

} // namespace ONNX_LIGHT_NAMESPACE::onnx_backend_test
