#include "stats/Uniform.hpp"

#include <numeric>

#include "utils/Invariant.hpp"

namespace ppforest2::stats {
  uint32_t Uniform::gen_lemire(uint32_t s, RNG& rng) const {
    uint32_t x = rng();
    uint64_t m = static_cast<uint64_t>(x) * s;
    uint32_t l = static_cast<uint32_t>(m);

    if (l < s) {
      uint32_t t = -s % s;

      while (l < t) {
        x = rng();
        m = static_cast<uint64_t>(x) * s;
        l = static_cast<uint32_t>(m);
      }
    }

    uint32_t result = m >> 32;

    invariant(result < s, "Uniform::gen_lemire: Invalid result generated by gen_lemire");

    return result;
  }

  Uniform::Uniform(int min, int max)
      : min(min)
      , max(max) {
    invariant(min >= 0, "Uniform: min must be greater than or equal to 0");
    invariant(min <= max, "Uniform: min must be less than or equal to max");
    invariant(max < std::numeric_limits<int>::max(), "Uniform: max must be less than the maximum value of an int");
  }

  int Uniform::operator()(RNG& rng) const {
    return min + static_cast<int>(gen_lemire(max + 1 - min, rng));
  }

  std::vector<int> Uniform::operator()(int count, RNG& rng) const {
    std::vector<int> result(count);

    for (int i = 0; i < count; i++) {
      result[i] = operator()(rng);
    }

    return result;
  }

  std::vector<int> Uniform::distinct(int count, RNG& rng) {
    int range_size = max - min + 1;

    invariant(count <= range_size, "Uniform::distinct: count must be cannot be greater than the range size");
    invariant(count >= 0, "Uniform::distinct: count must be greater than or equal to 0");

    std::vector<int> values(range_size);
    std::iota(values.begin(), values.end(), min);

    for (int i = range_size - 1; i >= 0; i--) {
      uint32_t j = gen_lemire(i + 1, rng);

      invariant(j < static_cast<uint32_t>(range_size), "Uniform::distinct: Invalid index generated by gen_lemire");

      std::swap(values[i], values[j]);
    }

    return std::vector<int>(values.begin(), values.begin() + count);
  }
}
