// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef CC_RANDOM_H_
#define CC_RANDOM_H_

#include <cstdint>
#include <random>

#include "absl/types/span.h"

namespace minigo {

// The C++ random library functionality is about as user friendly as its time
// library.
class Random {
 public:
  static constexpr uint64_t kLargePrime = 6364136223846793005ULL;
  static constexpr uint64_t kUniqueSeed = 0;
  static constexpr int kUniqueStream = 0;

  // The implementation supports generating multiple streams of uncorrelated
  // random numbers from a single seed.
  // If seed == Random::kUniqueSeed, a seed will be chosen from the platform's
  // random entropy source.
  // If stream == Random::kUniqueStream, a stream will be chosen from a
  // thread-safe global incrementing ID.
  // It's recommended that for reproducible results (modulo threading timing),
  // all Random instances use a seed speficied by a flag, and
  // Random::kUniqueStream for the stream.
  explicit Random(uint64_t seed, int stream);

  // Draw samples from a Dirichlet distribution.
  void Dirichlet(float alpha, absl::Span<float> samples);

  // Draw samples from a Dirichlet distribution.
  template <typename T>
  void Dirichlet(float alpha, T* array_like) {
    Dirichlet(alpha, {array_like->data(), array_like->size()});
  }

  // Draw samples from a Dirichlet distribution.
  template <int N>
  std::array<float, N> Dirichlet(float alpha) {
    std::array<float, N> samples;
    Dirichlet(alpha, &samples);
    return samples;
  }

  // Draw a single unform random samples in the half-open range [a, b).
  float Uniform(float a, float b);

  // Draw multiple unform random samples in the half-open range [a, b).
  void Uniform(float a, float b, absl::Span<float> samples);

  // Draw multiple unform random samples in the half-open range [a, b).
  template <typename T>
  void Uniform(float a, float b, T* array_like) {
    Uniform(a, b, {array_like->data(), array_like->size()});
  }

  // Draw multiple unform random samples in the half-open range [0, 1).
  template <typename T>
  void Uniform(T* array_like) {
    Uniform(0, 1, array_like);
  }

  // Draw multiple random samples from a normal distribution.
  template <typename T>
  void NormalDistribution(float mean, float stddev, T* array_like) {
    NormalDistribution(mean, stddev, {array_like->data(), array_like->size()});
  }

  // Draw multiple random samples from a normal distribution.
  void NormalDistribution(float mean, float stddev, absl::Span<float> samples);

  // Draw a single random sample from a normal distribution.
  float NormalDistribution(float mean, float stddev);

  // Returns a uniform random integer in the closed range [a, b].
  int UniformInt(int a, int b) {
    std::uniform_int_distribution<int> distribution(a, b);
    return distribution(impl_);
  }

  // Samples the given CDF at random, returning the index of the element found.
  // Guarantees that elements with zero probability
  int SampleCdf(absl::Span<float> cdf);

  uint64_t UniformUint64() {
    uint64_t a = impl_();
    uint64_t b = impl_();
    return (a << 32) | b;
  }

  // Returns a uniform random number in the half-open range [0, 1).
  float operator()() {
    return std::uniform_real_distribution<float>(0, 1)(impl_);
  }

  uint64_t state() const { return impl_.state; }
  uint64_t seed() const { return seed_; }
  int stream() const { return static_cast<int>(impl_.inc >> 1); }

  // Mixes the 64 bits into 32 bits that have improved entropy.
  // Useful if you have a 64 bit number with weaker entropy.
  static inline uint32_t MixBits(uint64_t x) {
    uint32_t xor_shifted = ((x >> 18u) ^ x) >> 27u;
    uint32_t rot = x >> 59u;
    return (xor_shifted >> rot) | (xor_shifted << ((-rot) & 31));
  }

  template <typename T>
  void Shuffle(T* array_like) {
    std::shuffle(array_like->begin(), array_like->end(), impl_);
  }

 private:
  // The implementation is based on 32bit PCG Random:
  //   http://www.pcg-random.org/
  struct Impl {
    using result_type = uint32_t;
    static constexpr result_type min() { return 0; }
    static constexpr result_type max() { return 0xffffffff; }

    Impl(uint64_t seed, int stream)
        : state(0), inc((static_cast<uint64_t>(stream) << 1) | 1) {
      operator()();
      state += seed;
      operator()();
    }

    result_type operator()() {
      auto result = MixBits(state);
      state = state * kLargePrime + inc;
      return result;
    }

    uint64_t state;
    const uint64_t inc;
  };

  uint64_t seed_;
  Impl impl_;
};

}  // namespace minigo

#endif  // CC_RANDOM_H_
