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

#include "onnx_core/backend_test/expect.h"
#include "onnx_extensions/backend_test/cases/training/include_training_cases.h"
#include "onnx_extensions/kernels/kernels/training/include_training_kernels.h"
#include "onnx_proto/onnx_helper.h"

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

namespace ONNX_LIGHT_NAMESPACE::onnx_backend_test {

namespace {

// See ``cases_adam.cc`` for the rationale on duplicating the domain
// constant rather than depending on ``lib_onnx_op``.
constexpr const char *kOnnxPreviewTrainingDomain = "ai.onnx.preview.training";

OpsetId TrainingOpset(int64_t version) { return OpsetId(kOnnxPreviewTrainingDomain, version); }

} // namespace

// ---------------------------------------------------------------------------
// Momentum — one iteration of the SGD-with-Momentum stochastic gradient
// optimization algorithm (since opset 1 in the
// ``ai.onnx.preview.training`` domain).
//
// Mirrors the upstream Python cases produced by
// ``onnx.backend.test.case.node.momentum.Momentum``:
//
//   * ``test_momentum``           — Single optimized rank-1 tensor,
//                                   ``mode == "standard"``.
//   * ``test_nesterov_momentum``  — Single optimized rank-1 tensor,
//                                   ``mode == "nesterov"``.
//   * ``test_momentum_multiple``  — Two optimized rank-1 tensors of
//                                   lengths 1 and 2, ``mode == "standard"``.
//
// All cases set the four attributes ``norm_coefficient``, ``alpha``,
// ``beta`` (FLOAT) and ``mode`` (STRING) and use the first-iteration
// path (``T == 0``).
// ---------------------------------------------------------------------------
void RegisterMomentumCases(std::vector<TestCase> &registry, TestMode mode) {
  const OpsetId opset = TrainingOpset(1);
  const KernelContext ctx{opset};
  const OpsetId default_opset = DefaultOpset(13);
  const onnx_kernels::kernel::Momentum momentum{ctx};

  if (mode == TestMode::BENCHMARK) {
    NodeProto node;
    node.set_op_type("Momentum");
    node.set_domain(kOnnxPreviewTrainingDomain);
    AddInputs(node, {"R", "T", "X", "G", "V"});
    AddOutputs(node, {"X_new", "V_new"});

    const float norm_coefficient = 0.001f;
    const float alpha = 0.95f;
    const float beta = 0.1f;
    AddFloatAttribute(node, "norm_coefficient", norm_coefficient);
    AddFloatAttribute(node, "alpha", alpha);
    AddFloatAttribute(node, "beta", beta);
    AddAttribute(node, "mode", std::string("standard"));

    Expect(registry, std::move(node), "test_momentum_benchmark", {default_opset, opset},
           {1, 1, kBenchmarkElementwiseSize, kBenchmarkElementwiseSize, kBenchmarkElementwiseSize},
           {kBenchmarkElementwiseSize, kBenchmarkElementwiseSize},
           [momentum, alpha, beta, norm_coefficient]() -> IoData {
             Tensor R = Tensor::FromFloat("", {}, {0.1f});
             Tensor T = Tensor::FromInt64("", {}, {0});
             Tensor X = RandnTensor(DataType::FLOAT, {kBenchmarkElementwiseSize}, 987654321ULL);
             Tensor G = RandnTensor(DataType::FLOAT, {kBenchmarkElementwiseSize}, 987654322ULL);
             Tensor V = RandnTensor(DataType::FLOAT, {kBenchmarkElementwiseSize}, 987654323ULL);
             std::vector<Tensor> outs = momentum(R, T, {X}, {G}, {V}, alpha, beta, norm_coefficient,
                                                 onnx_kernels::kernel::Momentum::Mode::kStandard);
             return IoData{{std::move(R), std::move(T), std::move(X), std::move(G), std::move(V)},
                           {std::move(outs[0]), std::move(outs[1])}};
           });
    return;
  }

  // From Momentum.export_momentum():
  {
    NodeProto node;
    node.set_op_type("Momentum");
    node.set_domain(kOnnxPreviewTrainingDomain);
    AddInputs(node, {"R", "T", "X", "G", "V"});
    AddOutputs(node, {"X_new", "V_new"});

    const float norm_coefficient = 0.001f;
    const float alpha = 0.95f;
    const float beta = 0.1f;
    AddFloatAttribute(node, "norm_coefficient", norm_coefficient);
    AddFloatAttribute(node, "alpha", alpha);
    AddFloatAttribute(node, "beta", beta);
    AddAttribute(node, "mode", std::string("standard"));
    Expect(registry, std::move(node), "test_momentum", {default_opset, opset}, [=]() -> IoData {
      Tensor R = Tensor::FromFloat("", {}, {0.1f});
      Tensor T = Tensor::FromInt64("", {}, {0});
      Tensor X = Tensor::FromFloat("", {2}, {1.2f, 2.8f});
      Tensor G = Tensor::FromFloat("", {2}, {-0.94f, -2.5f});
      Tensor V = Tensor::FromFloat("", {2}, {1.7f, 3.6f});

      std::vector<Tensor> outs = momentum(R, T, {X}, {G}, {V}, alpha, beta, norm_coefficient,
                                          onnx_kernels::kernel::Momentum::Mode::kStandard);
      return IoData{{std::move(R), std::move(T), std::move(X), std::move(G), std::move(V)},
                    {std::move(outs[0]), std::move(outs[1])}};
    });
  }

  // From Momentum.export_nesterov_momentum():
  {
    NodeProto node;
    node.set_op_type("Momentum");
    node.set_domain(kOnnxPreviewTrainingDomain);
    AddInputs(node, {"R", "T", "X", "G", "V"});
    AddOutputs(node, {"X_new", "V_new"});

    const float norm_coefficient = 0.01f;
    const float alpha = 0.95f;
    const float beta = 1.0f;
    AddFloatAttribute(node, "norm_coefficient", norm_coefficient);
    AddFloatAttribute(node, "alpha", alpha);
    AddFloatAttribute(node, "beta", beta);
    AddAttribute(node, "mode", std::string("nesterov"));
    Expect(registry, std::move(node), "test_nesterov_momentum", {default_opset, opset},
           [=]() -> IoData {
             Tensor R = Tensor::FromFloat("", {}, {0.1f});
             Tensor T = Tensor::FromInt64("", {}, {0});
             Tensor X = Tensor::FromFloat("", {2}, {1.2f, 2.8f});
             Tensor G = Tensor::FromFloat("", {2}, {-0.94f, -2.5f});
             Tensor V = Tensor::FromFloat("", {2}, {1.7f, 3.6f});

             std::vector<Tensor> outs = momentum(R, T, {X}, {G}, {V}, alpha, beta, norm_coefficient,
                                                 onnx_kernels::kernel::Momentum::Mode::kNesterov);
             return IoData{{std::move(R), std::move(T), std::move(X), std::move(G), std::move(V)},
                           {std::move(outs[0]), std::move(outs[1])}};
           });
  }

  // From Momentum.export_momentum_multiple():
  // NOTE: the upstream Python case labels the momentum-buffer inputs
  // "H1"/"H2" (an artefact of the test's source code), even though the
  // outputs use the "V*_new" naming. The string labels are purely cosmetic
  // — positionally these inputs are the per-tensor momentum buffers V1/V2
  // that the kernel reads. We mirror the upstream names verbatim so the
  // generated NodeProto matches the upstream ``test_momentum_multiple``
  // model byte-for-byte.
  {
    NodeProto node;
    node.set_op_type("Momentum");
    node.set_domain(kOnnxPreviewTrainingDomain);
    AddInputs(node, {"R", "T", "X1", "X2", "G1", "G2", "H1", "H2"});
    AddOutputs(node, {"X1_new", "X2_new", "V1_new", "V2_new"});

    const float norm_coefficient = 0.001f;
    const float alpha = 0.95f;
    const float beta = 0.85f;
    AddFloatAttribute(node, "norm_coefficient", norm_coefficient);
    AddFloatAttribute(node, "alpha", alpha);
    AddFloatAttribute(node, "beta", beta);
    AddAttribute(node, "mode", std::string("standard"));
    Expect(registry, std::move(node), "test_momentum_multiple", {default_opset, opset},
           [=]() -> IoData {
             Tensor R = Tensor::FromFloat("", {}, {0.1f});
             Tensor T = Tensor::FromInt64("", {}, {0});
             Tensor X1 = Tensor::FromFloat("", {1}, {1.0f});
             Tensor X2 = Tensor::FromFloat("", {2}, {1.0f, 2.0f});
             Tensor G1 = Tensor::FromFloat("", {1}, {-1.0f});
             Tensor G2 = Tensor::FromFloat("", {2}, {-1.0f, -3.0f});
             Tensor V1 = Tensor::FromFloat("", {1}, {2.0f});
             Tensor V2 = Tensor::FromFloat("", {2}, {4.0f, 1.0f});

             std::vector<Tensor> outs =
                 momentum(R, T, {X1, X2}, {G1, G2}, {V1, V2}, alpha, beta, norm_coefficient,
                          onnx_kernels::kernel::Momentum::Mode::kStandard);
             return IoData{
                 {std::move(R), std::move(T), std::move(X1), std::move(X2), std::move(G1),
                  std::move(G2), std::move(V1), std::move(V2)},
                 {std::move(outs[0]), std::move(outs[1]), std::move(outs[2]), std::move(outs[3])}};
           });
  }
}

} // namespace ONNX_LIGHT_NAMESPACE::onnx_backend_test
