#pragma once

#include <psemek/random/uniform.hpp>

#include <vector>
#include <algorithm>

namespace psemek::random
{

	// TODO: reimplement in O(1) using Vose's Alias Algorithm
	// https://www.keithschwarz.com/darts-dice-coins/

	template <typename T>
	struct weighted_distribution
	{
		weighted_distribution() = default;

		weighted_distribution(std::vector<T> frequencies)
			: prefix_sum_(std::move(frequencies))
		{
			T sum = T{0};
			for (auto const & w : prefix_sum_)
				sum += w;

			for (auto & w : prefix_sum_)
				w /= sum;

			for (std::size_t i = 1; i < prefix_sum_.size(); ++i)
				prefix_sum_[i] += prefix_sum_[i - 1];
		}

		std::size_t size() const { return prefix_sum_.size(); }

		template <typename RNG>
		std::size_t operator()(RNG && rng) const
		{
			auto roll = uniform<T>(rng);
			auto it = std::upper_bound(prefix_sum_.begin(), prefix_sum_.end(), roll);
			if (it == prefix_sum_.end()) --it;
			return it - prefix_sum_.begin();
		}

		std::vector<T> release()
		{
			return std::move(prefix_sum_);
		}

	private:
		std::vector<T> prefix_sum_;
	};

}
