#pragma once

#include <psemek/math/interval.hpp>

#include <limits>
#include <type_traits>

namespace psemek::random
{

	template <typename T>
	struct uniform_int_distribution
	{
		static_assert(std::is_integral_v<T>);

		using result_type = T;

		uniform_int_distribution()
			: range_{std::numeric_limits<T>::min(), std::numeric_limits<T>::max()}
		{}

		uniform_int_distribution(T min, T max)
			: range_{min, max}
		{}

		uniform_int_distribution(math::interval<T> range)
			: range_(range)
		{}

		uniform_int_distribution(uniform_int_distribution const &) = default;

		template <typename RNG>
		T operator()(RNG && rng);

		math::interval<T> range() const { return range_; }

	private:
		math::interval<T> range_;

		template <typename RNG>
		static T generate(RNG & rng, math::interval<T> const & range);
	};

	template <typename T>
	template <typename RNG>
	T uniform_int_distribution<T>::operator()(RNG && rng)
	{
		return generate(rng, range_);
	}

	template <typename T>
	template <typename RNG>
	T uniform_int_distribution<T>::generate(RNG & rng, math::interval<T> const & range)
	{
		using ctype = std::common_type_t<typename RNG::result_type, std::make_unsigned_t<T>>;

		ctype const rng_min = rng.min();
		ctype const rng_max = rng.max();
		ctype const rng_range = rng_max - rng_min;
		ctype const gen_range = static_cast<ctype>(range.max) - static_cast<ctype>(range.min);

		ctype result;

		if (rng_range == gen_range)
		{
			result = static_cast<ctype>(rng()) - rng_min;
		}
		else if (rng_range > gen_range)
		{
			ctype const gen_range_full = gen_range + 1;
			ctype const buckets = rng_range / gen_range_full;
			ctype const max = gen_range_full * buckets;

			do
			{
				result = static_cast<ctype>(rng()) - rng_min;
			} while (result >= max);
			result /= buckets;
		}
		else // if (rng_range < gen_range)
		{
			ctype tmp;
			do
			{
				const ctype rng_range_full = rng_range + 1;
				tmp = rng_range_full * generate(rng, {0, gen_range / rng_range_full});
				result = tmp + (static_cast<ctype>(rng()) - rng_min);
			}
			while (result > gen_range || result < tmp);
		}

		return result + range.min;
	}

}
