#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
// Philox Counter based RNG implementation for Metal
// Borrowed from aten/src/ATen/core/PhiloxRNGEngine.h
// Which in turn borrowed from
// http://www.thesalmons.org/john/random123/papers/random123sc11.pdf
#pragma once
#include <metal_stdlib>

namespace c10 {
namespace metal {

namespace detail {

constexpr float uint32_to_uniform_float(uint32_t value) {
  // maximum value such that `MAX_INT * scale < 1.0` (with float rounding)
  constexpr float scale = 4.6566127342e-10;
  return static_cast<float>(value & 0x7FFFFFFF) * scale;
}

constexpr uint2 splitlong(ulong v) {
  return uint2(v >> 32, v & 0xffffffff);
}

} // namespace detail

namespace philox4 {

constexpr uint2 mulhilo(uint a, uint b) {
  auto rc = static_cast<ulong>(a) * b;
  return detail::splitlong(rc);
}
inline uint4 single_round(uint4 ctr, uint2 key) {
  constexpr uint kPhiloxSA = 0xD2511F53;
  constexpr uint kPhiloxSB = 0xCD9E8D57;
  auto rc0 = mulhilo(kPhiloxSA, ctr.x);
  auto rc1 = mulhilo(kPhiloxSB, ctr.z);
  return uint4(rc1.x ^ ctr.y ^ key.x, rc1.y, rc0.x ^ ctr.w ^ key.y, rc0.y);
}

inline uint4 multiple_rounds(uint4 ctr, uint2 key, uint rounds) {
  constexpr uint2 kPhilox10 = {0x9E3779B9, 0xBB67AE85};
  for (uint round = 0; round < rounds - 1; ++round) {
    ctr = single_round(ctr, key);
    key += kPhilox10;
  }
  return ctr;
}

inline uint4 rand(long seed, long index) {
  uint4 ctr = 0;
  ctr.zw = detail::splitlong(index);
  return multiple_rounds(ctr, detail::splitlong(seed), 10);
}

} // namespace philox4

inline float randn(long seed, long index) {
  auto value = philox4::rand(seed, index);
  float u1 = 1.0 - detail::uint32_to_uniform_float(value.x);
  float u2 = 1.0 - detail::uint32_to_uniform_float(value.y);
  return ::metal::sqrt(-2.0 * ::metal::log(u1)) *
      ::metal::cos(2.0 * M_PI_F * u2);
}

// Box-Muller transform: convert two raw Philox uint32 outputs into two
// independent N(0, 1) samples. The radial uniform `u1` is clamped away from
// `0` so `log(u1)` stays finite; the angular uniform `u2` needs no clamp
// because `sin` / `cos` are bounded over the entire domain. `sincos` is the
// fused intrinsic that returns `sin` and writes `cos` in one call.
inline float2 box_muller_from_philox(uint2 raw) {
  constexpr float eps = ::metal::numeric_limits<float>::epsilon();
  float u1 = ::metal::max(detail::uint32_to_uniform_float(raw.x), eps);
  float u2 = detail::uint32_to_uniform_float(raw.y);
  float r = ::metal::precise::sqrt(-2.0f * ::metal::precise::log(u1));
  float c;
  float s = ::metal::precise::sincos(2.0f * M_PI_F * u2, c);
  return r * float2(c, s);
}

inline float rand(long seed, long index) {
  auto value = philox4::rand(seed, index);
  return detail::uint32_to_uniform_float(value.x);
}

inline long randint64(long seed, long index, long low, long high) {
  auto range = high - low;
  auto value = philox4::rand(seed, index);
  // TODO: Implement better algorithm for large ranges
  return low +
      static_cast<long>(detail::uint32_to_uniform_float(value.x) * range);
}

} // namespace metal
} // namespace c10

#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)