#pragma once

#include <psemek/math/vector.hpp>
#include <psemek/math/point.hpp>
#include <psemek/math/constants.hpp>
#include <psemek/random/normal.hpp>
#include <psemek/random/uniform_int.hpp>

namespace psemek::random
{

	template <typename T, std::size_t N>
	struct uniform_sphere_vector_distribution
	{
		using result_type = math::vector<T, N>;

		uniform_sphere_vector_distribution(T r = T{1})
			: r_{r}
		{}

		uniform_sphere_vector_distribution(uniform_sphere_vector_distribution const &) = default;

		template <typename RNG>
		auto operator()(RNG && rng)
		{
			result_type v;
			while (true)
			{
				for (std::size_t i = 0; i < N; ++i)
					v[i] = d_(rng);

				T l = length(v);
				if (l != T{})
					return v / l * r_;
			}
		}

		T radius() const { return r_; }

	private:
		normal_distribution<T> d_;
		T r_;
	};

	template <typename T>
	struct uniform_sphere_vector_distribution<T, 1>
	{
		using result_type = math::vector<T, 1>;

		uniform_sphere_vector_distribution(T r = T{1})
			: r_{r}
		{}

		uniform_sphere_vector_distribution(uniform_sphere_vector_distribution const &) = default;

		template <typename RNG>
		math::vector<T, 1> operator()(RNG && rng)
		{
			uniform_int_distribution<int> d{0, 1};
			if (d(rng) == 0)
				return {r_};
			return {-r_};
		}

		T radius() const { return r_; }

	private:
		T r_;
	};

	template <typename T>
	struct uniform_sphere_vector_distribution<T, 2>
	{
		using result_type = math::vector<T, 2>;

		uniform_sphere_vector_distribution(T r = T{1})
			: r_{r}
		{}

		uniform_sphere_vector_distribution(uniform_sphere_vector_distribution const &) = default;

		template <typename RNG>
		result_type operator()(RNG && rng)
		{
			uniform_real_distribution<T> d{0, 2 * math::pi};
			T a = d(rng);
			return r_ * result_type{std::cos(a), std::sin(a)};
		}

		T radius() const { return r_; }

	private:
		T r_;
	};

	template <typename T>
	struct uniform_sphere_vector_distribution<T, 3>
	{
		using result_type = math::vector<T, 3>;

		uniform_sphere_vector_distribution(T r = T{1})
			: r_{r}
		{}

		template <typename RNG>
		auto operator()(RNG && rng)
		{
			uniform_real_distribution<T> d{-1, 1};

			while (true)
			{
				T const x = d(rng);
				T const y = d(rng);

				T const s = x * x + y * y;

				if (s > 1) continue;

				T const m = 2 * std::sqrt(1 - s);

				return r_ * result_type{x * m, y * m, 2 * s - 1};
			}
		}

		T radius() const { return r_; }

	private:
		T r_;
	};

	template <typename T, std::size_t N>
	struct uniform_sphere_point_distribution
	{
		using result_type = math::point<T, N>;

		uniform_sphere_point_distribution(result_type origin, T r = T{1})
			: origin_{origin}
			, vector_d_{r}
		{}

		template <typename RNG>
		auto operator()(RNG && rng)
		{
			return origin_ + vector_d_(rng);
		}

		result_type origin() const { return origin_; }

		T radius() const { return vector_d_.radius(); }

	private:
		result_type origin_;
		uniform_sphere_vector_distribution<T, N> vector_d_;
	};


}
