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

#include "onnx_core/runtime/kernels/random.h"

#include "onnx_core/compute/raw_buffer_allocator.h"

#include <gtest/gtest.h>

#include <array>
#include <cstring>
#include <vector>

using namespace ONNX_LIGHT_NAMESPACE;

namespace Test {

TEST(BackendRandomTest, RandIsDeterministicWithSeed) {
  auto first = core::runtime::Rand({2, 3}, /*seed=*/42);
  auto second = core::runtime::Rand({2, 3}, /*seed=*/42);
  EXPECT_EQ(first, second);
  EXPECT_EQ(first.size(), 6u);
}

TEST(BackendRandomTest, RandDefaultSeedMatchesPythonReference) {
  // Values produced by the Python reference for ``rand(3, seed=7)``.
  std::vector<double> expected = {
      0.3898297483912715,
      0.01678829452815611,
      0.9007606806068834,
  };
  auto values = core::runtime::Rand({3}, /*seed=*/7);
  ASSERT_EQ(values.size(), expected.size());
  for (size_t i = 0; i < expected.size(); ++i) {
    EXPECT_DOUBLE_EQ(values[i], expected[i]);
  }
}

TEST(BackendRandomTest, RandnDefaultSeedMatchesPythonReference) {
  std::vector<double> expected = {
      -1.0009541026316917,
      1.3335403660869787,
      0.5188620948808316,
  };
  auto values = core::runtime::Randn({3}, /*seed=*/7);
  ASSERT_EQ(values.size(), expected.size());
  for (size_t i = 0; i < expected.size(); ++i) {
    EXPECT_DOUBLE_EQ(values[i], expected[i]);
  }
}

TEST(BackendRandomTest, RandnFloatMatchesCastFromDouble) {
  auto doubles = core::runtime::Randn<double>({5}, /*seed=*/7);
  auto floats = core::runtime::Randn<float>({5}, /*seed=*/7);
  ASSERT_EQ(doubles.size(), floats.size());
  for (size_t i = 0; i < doubles.size(); ++i) {
    EXPECT_FLOAT_EQ(floats[i], static_cast<float>(doubles[i]));
  }
}

TEST(BackendRandomTest, RandIntMatchesPythonReference) {
  std::vector<int64_t> expected = {2, 3, 0, 3, 5};
  auto values = core::runtime::RandInt(0, 7, {5}, /*seed=*/7);
  EXPECT_EQ(values, expected);
}

TEST(BackendRandomTest, RandIntBoundsAreRespected) {
  auto values = core::runtime::RandInt(0, 5, {100}, /*seed=*/12);
  ASSERT_EQ(values.size(), 100u);
  for (int64_t v : values) {
    EXPECT_GE(v, 0);
    EXPECT_LT(v, 5);
  }
}

TEST(BackendRandomTest, RandIntRejectsInvalidRange) {
  EXPECT_THROW(core::runtime::RandInt(5, 5, {3}), std::invalid_argument);
  EXPECT_THROW(core::runtime::RandInt(10, 2, {3}), std::invalid_argument);
}

TEST(BackendRandomTest, NegativeShapeRejected) {
  EXPECT_THROW(core::runtime::Rand({-1, 2}), std::invalid_argument);
  EXPECT_THROW(core::runtime::Randn({2, -3}), std::invalid_argument);
  EXPECT_THROW(core::runtime::RandInt(0, 2, {-1}), std::invalid_argument);
}

TEST(BackendRandomTest, EmptyShapeReturnsSingleValue) {
  EXPECT_EQ(core::runtime::Rand({}).size(), 1u);
  EXPECT_EQ(core::runtime::Randn({}).size(), 1u);
}

TEST(BackendRandomTest, RandnTensorSupportsNumericElementTypes) {
  constexpr std::array<int32_t, 9> element_types = {
      core::runtime::DataType::FLOAT,   core::runtime::DataType::DOUBLE,
      core::runtime::DataType::FLOAT16, core::runtime::DataType::BFLOAT16,
      core::runtime::DataType::INT8,    core::runtime::DataType::INT16,
      core::runtime::DataType::INT32,   core::runtime::DataType::INT64,
      core::runtime::DataType::BOOL,
  };
  for (int32_t element_type : element_types) {
    const core::runtime::Tensor first =
        core::runtime::RandnTensor(element_type, {2, 3}, /*seed=*/7);
    const core::runtime::Tensor second =
        core::runtime::RandnTensor(element_type, {2, 3}, /*seed=*/7);
    EXPECT_EQ(first.data_type, element_type);
    EXPECT_EQ(first.shape, (core::runtime::Shape{2, 3}));
    ASSERT_EQ(first.size_bytes(), second.size_bytes());
    EXPECT_EQ(std::memcmp(first.bytes(), second.bytes(), first.size_bytes()), 0);
  }
}

TEST(BackendRandomTest, RandnTensorUsesAllocator) {
  core::runtime::SimpleRawBufferAllocator allocator(1);
  const core::runtime::Tensor tensor =
      core::runtime::RandnTensor(core::runtime::DataType::FLOAT16, {4}, /*seed=*/7, &allocator);
  EXPECT_TRUE(tensor.has_allocation());
  EXPECT_EQ(tensor.allocation_owner(), &allocator);
}

TEST(BackendRandomTest, RandnTensorRejectsUnsupportedElementType) {
  EXPECT_THROW(core::runtime::RandnTensor(core::runtime::DataType::STRING, {2}),
               std::invalid_argument);
}

} // namespace Test
