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

#include "onnx_core/backend_test/expect.h"
#include "onnx_extensions/backend_test/cases/generator/include_generator_cases.h"
#include "onnx_extensions/kernels/kernels/generator/include_generator_kernels.h"

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

namespace ONNX_LIGHT_NAMESPACE::onnx_backend_test {

namespace {

// Builds a single-node ``Bernoulli`` topology with the given ``input``
// probability tensor and optional ``seed`` / ``dtype`` attributes, runs
// the reference kernel to obtain the expected output, and registers the
// case under ``case_name``.
//
// When ``seed`` is :cpp:var:`kernel::Bernoulli::kNoSeed`, no ``seed``
// attribute is added to the node. When ``dtype`` is ``0``, no ``dtype``
// attribute is added; otherwise the attribute holds the value passed and
// the kernel uses it as the output element type.
void RegisterOneBernoulli(const std::string &case_name, const Tensor &input, int64_t seed,
                          int32_t dtype, std::vector<TestCase> &registry, const OpsetId &opset) {
  NodeProto node;
  node.set_op_type("Bernoulli");
  node.add_input("x");
  node.add_output("y");

  if (seed != onnx_kernels::kernel::Bernoulli::kNoSeed) {
    AttributeProto *attr = node.add_attribute();
    attr->set_name("seed");
    attr->set_type(AttributeProto::AttributeType::FLOAT);
    attr->set_f(static_cast<float>(seed));
  }
  if (dtype != 0) {
    AttributeProto *attr = node.add_attribute();
    attr->set_name("dtype");
    attr->set_type(AttributeProto::AttributeType::INT);
    attr->set_i(static_cast<int64_t>(dtype));
  }
  Expect(registry, std::move(node), case_name, {opset}, [=]() -> IoData {
    const KernelContext ctx{opset};
    Tensor input_named = input;
    input_named.name = "x";
    Tensor y = onnx_kernels::kernel::Bernoulli(ctx)(input_named, seed, dtype);

    return IoData{{std::move(input_named)}, {std::move(y)}};
  });
}

} // namespace

// ---------------------------------------------------------------------------
// Bernoulli — draws binary samples from a Bernoulli distribution per
// element of the input probability tensor. The operator is officially
// non-deterministic, but the kernel uses a fixed seed by default so that
// these cases produce stable expected outputs.
//
// The cases mirror the upstream
// ``onnx.backend.test.case.node.bernoulli`` exports
// (``test_bernoulli``, ``test_bernoulli_double``,
// ``test_bernoulli_seed``) at the structural level while using outputs
// generated by this kernel so that the reference and the test fixture
// always agree.
// ---------------------------------------------------------------------------
void RegisterBernoulliCases(std::vector<TestCase> &registry, TestMode mode) {
  const OpsetId opset = DefaultOpset(22);

  const std::vector<float> probs = {0.0f, 0.1f, 0.25f, 0.4f, 0.5f, 0.6f, 0.75f, 0.9f, 1.0f, 0.5f};

  if (mode == TestMode::BENCHMARK) {
    NodeProto node;
    node.set_op_type("Bernoulli");
    node.add_input("x");
    node.add_output("y");

    const KernelContext ctx{opset};
    const onnx_kernels::kernel::Bernoulli bernoulli_kernel{ctx};
    Expect(registry, std::move(node), "test_cc_bernoulli_benchmark", {opset},
           {kBenchmarkElementwiseSize}, {kBenchmarkElementwiseSize},
           [bernoulli_kernel]() -> IoData {
             Tensor x = Tensor::FromFloat("x", /*shape=*/{kBenchmarkElementwiseSize},
                                          std::vector<float>(kBenchmarkElementwiseSize, 0.5f));
             Tensor y = bernoulli_kernel(x, onnx_kernels::kernel::Bernoulli::kNoSeed, /*dtype=*/0);
             return IoData{{std::move(x)}, {std::move(y)}};
           });
    return;
  }

  // Default attributes: no seed, dtype matches input (FLOAT).
  {
    const Tensor x = Tensor::FromFloat("", /*shape=*/{static_cast<int64_t>(probs.size())}, probs);
    RegisterOneBernoulli("test_cc_bernoulli", x, onnx_kernels::kernel::Bernoulli::kNoSeed,
                         /*dtype=*/0, registry, opset);
  }

  // dtype = DOUBLE: output element type is overridden by the attribute.
  {
    const Tensor x = Tensor::FromFloat("", /*shape=*/{static_cast<int64_t>(probs.size())}, probs);
    RegisterOneBernoulli("test_cc_bernoulli_double", x, onnx_kernels::kernel::Bernoulli::kNoSeed,
                         /*dtype=*/static_cast<int32_t>(DataType::DOUBLE), registry, opset);
  }

  // seed = 0: explicit ``seed`` attribute.
  {
    const Tensor x = Tensor::FromFloat("", /*shape=*/{static_cast<int64_t>(probs.size())}, probs);
    RegisterOneBernoulli("test_cc_bernoulli_seed", x, /*seed=*/0, /*dtype=*/0, registry, opset);
  }
}

} // namespace ONNX_LIGHT_NAMESPACE::onnx_backend_test
