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

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

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

#include <cstdint>
#include <limits>
#include <stdexcept>
#include <type_traits>

namespace ONNX_LIGHT_NAMESPACE::core::runtime {

namespace {

constexpr double kInvTwoPow53 = 1.0 / static_cast<double>(static_cast<uint64_t>(1) << 53);

int64_t ShapeToCount(const Shape &shape) {
  if (shape.empty()) {
    return 1;
  }
  int64_t count = 1;
  for (int64_t dim : shape) {
    EXT_ENFORCE_INVALID(dim >= 0, "shape cannot contain negative dimensions.");
    count *= dim;
  }
  return count;
}

inline double UniformFromState(uint64_t value) {
  return static_cast<double>(value >> 11) * kInvTwoPow53;
}

} // namespace

std::pair<uint64_t, uint64_t> NextUint64(uint64_t state) {
  state += 0x9E3779B97F4A7C15ULL;
  uint64_t mixed = state;
  mixed = (mixed ^ (mixed >> 30)) * 0xBF58476D1CE4E5B9ULL;
  mixed = (mixed ^ (mixed >> 27)) * 0x94D049BB133111EBULL;
  mixed = mixed ^ (mixed >> 31);
  return {state, mixed};
}

template <typename T> std::vector<T> Rand(const Shape &shape, std::optional<uint64_t> seed) {
  static_assert(std::is_floating_point_v<T>, "Rand<T> requires a floating-point element type.");
  const int64_t count = ShapeToCount(shape);
  std::vector<T> values(static_cast<size_t>(count));
  uint64_t state = seed.value_or(kDefaultSeed);
  for (int64_t i = 0; i < count; ++i) {
    auto [next_state, value] = NextUint64(state);
    state = next_state;
    values[static_cast<size_t>(i)] = static_cast<T>(UniformFromState(value));
  }
  return values;
}

template std::vector<double> Rand<double>(const Shape &shape, std::optional<uint64_t> seed);
template std::vector<float> Rand<float>(const Shape &shape, std::optional<uint64_t> seed);

std::vector<int64_t> RandInt(int64_t low, int64_t high, const Shape &shape,
                             std::optional<uint64_t> seed) {
  EXT_ENFORCE_INVALID(high > low, "high must be greater than low.");
  const int64_t count = ShapeToCount(shape);
  const uint64_t span = static_cast<uint64_t>(high - low);
  // Mirror the Python reference:
  //   limit = 2**64 - (2**64 % span)
  // Using unsigned 64-bit arithmetic: (0 - span) % span == 2**64 % span.
  const uint64_t modulo = (static_cast<uint64_t>(0) - span) % span;
  const uint64_t limit = static_cast<uint64_t>(0) - modulo;
  std::vector<int64_t> values(static_cast<size_t>(count));
  uint64_t state = seed.value_or(kDefaultSeed);
  for (int64_t i = 0; i < count; ++i) {
    uint64_t candidate;
    while (true) {
      auto [next_state, value] = NextUint64(state);
      state = next_state;
      // ``limit`` equals 0 when span is a power of two that divides 2**64; in
      // that case every candidate is acceptable.
      if (limit == 0 || value < limit) {
        candidate = value % span;
        break;
      }
    }
    values[static_cast<size_t>(i)] = static_cast<int64_t>(candidate) + low;
  }
  return values;
}

template <typename T> std::vector<T> Randn(const Shape &shape, std::optional<uint64_t> seed) {
  static_assert(std::is_floating_point_v<T>, "Randn<T> requires a floating-point element type.");
  const int64_t count = ShapeToCount(shape);
  std::vector<T> values(static_cast<size_t>(count));
  uint64_t state = seed.value_or(kDefaultSeed);
  for (int64_t i = 0; i < count; ++i) {
    double sample = 0.0;
    for (int j = 0; j < 12; ++j) {
      auto [next_state, value] = NextUint64(state);
      state = next_state;
      sample += UniformFromState(value);
    }
    values[static_cast<size_t>(i)] = static_cast<T>(sample - 6.0);
  }
  return values;
}

template std::vector<double> Randn<double>(const Shape &shape, std::optional<uint64_t> seed);
template std::vector<float> Randn<float>(const Shape &shape, std::optional<uint64_t> seed);

Tensor RandnTensor(int32_t element_type, const Shape &shape, std::optional<uint64_t> seed,
                   RawBufferAllocator *allocator) {
  const uint64_t resolved_seed = seed.value_or(kDefaultSeed);
  switch (static_cast<DataType>(element_type)) {
  case DataType::FLOAT:
    return Tensor::FromFloat("", shape, Randn<float>(shape, resolved_seed), allocator);
  case DataType::DOUBLE:
    return Tensor::FromDouble("", shape, Randn<double>(shape, resolved_seed), allocator);
  case DataType::FLOAT16:
    return MakeFloat16Tensor("", shape, Randn<float>(shape, resolved_seed), allocator);
  case DataType::BFLOAT16:
    return MakeBfloat16Tensor("", shape, Randn<float>(shape, resolved_seed), allocator);
  case DataType::INT8:
    return Tensor::FromInt8("", shape, RandnInt<int8_t>(shape, resolved_seed), allocator);
  case DataType::INT16:
    return Tensor::FromInt16("", shape, RandnInt<int16_t>(shape, resolved_seed), allocator);
  case DataType::INT32:
    return Tensor::FromInt32("", shape, RandnInt<int32_t>(shape, resolved_seed), allocator);
  case DataType::INT64:
    return Tensor::FromInt64("", shape, RandnInt<int64_t>(shape, resolved_seed), allocator);
  case DataType::BOOL:
    return RandBool(shape, resolved_seed, allocator);
  default:
    throw std::invalid_argument("RandnTensor does not support element type " +
                                std::to_string(element_type) + ".");
  }
}

template <typename T>
void RandUniformInto(T *dst, int64_t count, double low, double high, std::optional<uint64_t> seed) {
  static_assert(std::is_floating_point_v<T>,
                "RandUniformInto<T> requires a floating-point element type.");
  const T span = static_cast<T>(high - low);
  const T lowT = static_cast<T>(low);
  uint64_t state = seed.value_or(kDefaultSeed);
  for (int64_t i = 0; i < count; ++i) {
    auto [next_state, value] = NextUint64(state);
    state = next_state;
    dst[static_cast<size_t>(i)] = lowT + static_cast<T>(UniformFromState(value)) * span;
  }
}

template void RandUniformInto<double>(double *dst, int64_t count, double low, double high,
                                      std::optional<uint64_t> seed);
template void RandUniformInto<float>(float *dst, int64_t count, double low, double high,
                                     std::optional<uint64_t> seed);

template <typename T>
void RandNormalInto(T *dst, int64_t count, double mean, double scale,
                    std::optional<uint64_t> seed) {
  static_assert(std::is_floating_point_v<T>,
                "RandNormalInto<T> requires a floating-point element type.");
  const T meanT = static_cast<T>(mean);
  const T scaleT = static_cast<T>(scale);
  uint64_t state = seed.value_or(kDefaultSeed);
  for (int64_t i = 0; i < count; ++i) {
    double sample = 0.0;
    for (int j = 0; j < 12; ++j) {
      auto [next_state, value] = NextUint64(state);
      state = next_state;
      sample += UniformFromState(value);
    }
    dst[static_cast<size_t>(i)] = meanT + static_cast<T>(sample - 6.0) * scaleT;
  }
}

template void RandNormalInto<double>(double *dst, int64_t count, double mean, double scale,
                                     std::optional<uint64_t> seed);
template void RandNormalInto<float>(float *dst, int64_t count, double mean, double scale,
                                    std::optional<uint64_t> seed);

template <typename TInt> std::vector<TInt> RandnInt(const Shape &shape, uint64_t seed) {
  const std::vector<float> floats = Randn<float>(shape, seed);
  std::vector<TInt> out(floats.size());
  for (size_t i = 0; i < floats.size(); ++i) {
    out[i] = static_cast<TInt>(floats[i]);
  }
  return out;
}

template std::vector<int8_t> RandnInt<int8_t>(const Shape &shape, uint64_t seed);
template std::vector<int16_t> RandnInt<int16_t>(const Shape &shape, uint64_t seed);
template std::vector<int32_t> RandnInt<int32_t>(const Shape &shape, uint64_t seed);
template std::vector<int64_t> RandnInt<int64_t>(const Shape &shape, uint64_t seed);

template <typename TUInt>
std::vector<TUInt> RandUint(int64_t high, const Shape &shape, uint64_t seed) {
  const std::vector<int64_t> ints = RandInt(0, high, shape, seed);
  std::vector<TUInt> out(ints.size());
  for (size_t i = 0; i < ints.size(); ++i) {
    out[i] = static_cast<TUInt>(ints[i]);
  }
  return out;
}

template std::vector<uint8_t> RandUint<uint8_t>(int64_t high, const Shape &shape, uint64_t seed);
template std::vector<uint16_t> RandUint<uint16_t>(int64_t high, const Shape &shape, uint64_t seed);
template std::vector<uint32_t> RandUint<uint32_t>(int64_t high, const Shape &shape, uint64_t seed);
template std::vector<uint64_t> RandUint<uint64_t>(int64_t high, const Shape &shape, uint64_t seed);

Tensor RandBool(const Shape &shape, std::optional<uint64_t> seed, RawBufferAllocator *allocator) {
  const std::vector<double> values = Randn<double>(shape, seed);
  Tensor tensor = MakeOutputTensor(DataType::BOOL, shape, values.size(), allocator);
  uint8_t *bytes = tensor.mutable_bytes();
  for (size_t i = 0; i < values.size(); ++i) {
    bytes[i] = values[i] > 0.0 ? 1 : 0;
  }
  return tensor;
}

} // namespace ONNX_LIGHT_NAMESPACE::core::runtime
