#pragma once

#include <BAN/Heap.h>
#include <BAN/Math.h>
#include <BAN/Swap.h>
#include <BAN/Traits.h>
#include <BAN/Vector.h>

namespace BAN::sort
{

	template<typename It, typename Comp = less<it_value_type_t<It>>>
	void exchange_sort(It begin, It end, Comp comp = {})
	{
		for (It lhs = begin; lhs != end; ++lhs)
			for (It rhs = next(lhs, 1); rhs != end; ++rhs)
				if (!comp(*lhs, *rhs))
					swap(*lhs, *rhs);
	}

	namespace detail
	{

		template<typename It>
		struct partition_pair
		{
			It lt;
			It gt;
		};

		template<typename It, typename Comp>
		partition_pair<It> partition(It begin, It end, Comp comp)
		{
			It pivot = next(begin, distance(begin, end) / 2);

			It lt = begin;
			It eq = begin;
			It gt = end;

			while (eq != gt)
			{
				if (comp(*eq, *pivot))
				{
					swap(*eq, *lt);
					if (pivot == lt)
						pivot = eq;
					++lt;
					++eq;
				}
				else if (comp(*pivot, *eq))
				{
					--gt;
					swap(*eq, *gt);
					if (pivot == gt)
						pivot = eq;
				}
				else
				{
					++eq;
				}
			}

			return { lt, gt };
		}
	}

	template<typename It, typename Comp = less<it_value_type_t<It>>>
	void quick_sort(It begin, It end, Comp comp = {})
	{
		if (distance(begin, end) <= 1)
			return;
		const auto [lt, gt] = detail::partition(begin, end, comp);
		quick_sort(begin, lt, comp);
		quick_sort(gt, end, comp);
	}

	template<typename It, typename Comp = less<it_value_type_t<It>>>
	void insertion_sort(It begin, It end, Comp comp = {})
	{
		if (distance(begin, end) <= 1)
			return;
		for (It it1 = next(begin, 1); it1 != end; ++it1)
		{
			auto x = move(*it1);
			It it2 = it1;
			for (; it2 != begin && comp(x, *prev(it2, 1)); --it2)
				*it2 = move(*prev(it2, 1));
			*it2 = move(x);
		}
	}

	template<typename It, typename Comp = less<it_value_type_t<It>>>
	void heap_sort(It begin, It end, Comp comp = {})
	{
		make_heap(begin, end, comp);
		sort_heap(begin, end, comp);
	}

	namespace detail
	{

		template<typename It, typename Comp>
		void intro_sort_impl(It begin, It end, size_t max_depth, Comp comp)
		{
			if (distance(begin, end) <= 16)
				return insertion_sort(begin, end, comp);
			if (max_depth == 0)
				return heap_sort(begin, end, comp);
			const auto [lt, gt] = detail::partition(begin, end, comp);
			intro_sort_impl(begin, lt, max_depth - 1, comp);
			intro_sort_impl(gt, end, max_depth - 1, comp);
		}

	}

	template<typename It, typename Comp = less<it_value_type_t<It>>>
	void intro_sort(It begin, It end, Comp comp = {})
	{
		const size_t len = distance(begin, end);
		if (len <= 1)
			return;
		detail::intro_sort_impl(begin, end, 2 * Math::ilog2(len), comp);
	}

	namespace detail
	{

		template<unsigned_integral T>
		consteval T lsb_index(T value)
		{
			for (T result = 0;; result++)
				if (value & (1 << result))
					return result;
		}

	}

	template<typename It, size_t radix = 256>
	requires is_unsigned_v<it_value_type_t<It>> && (radix > 0 && (radix & (radix - 1)) == 0)
	void radix_sort(It begin, It end, BAN::Span<it_value_type_t<It>> storage)
	{
		const size_t len = distance(begin, end);
		if (len <= 1)
			return;

		ASSERT(storage.size() >= len);

		constexpr size_t mask  = radix - 1;
		constexpr size_t shift = detail::lsb_index(radix);

		for (size_t s = 0; s < sizeof(it_value_type_t<It>) * 8; s += shift)
		{
			size_t counts[radix] {};
			for (It it = begin; it != end; ++it)
				counts[(*it >> s) & mask]++;

			for (size_t i = 0; i < radix - 1; i++)
				counts[i + 1] += counts[i];

			for (It it = end; it != begin;)
			{
				--it;
				storage[--counts[(*it >> s) & mask]] = *it;
			}

			It it = begin;
			for (size_t j = 0; j < storage.size(); j++, ++it)
				*it = storage[j];
		}
	}

	template<typename It, size_t radix = 256>
	requires is_unsigned_v<it_value_type_t<It>> && (radix > 0 && (radix & (radix - 1)) == 0)
	BAN::ErrorOr<void> radix_sort(It begin, It end)
	{
		const size_t len = distance(begin, end);
		if (len <= 1)
			return {};

		Vector<it_value_type_t<It>> temp;
		TRY(temp.resize(len));

		radix_sort(begin, end, temp.span());

		return {};
	}

	template<typename It, typename Comp = less<it_value_type_t<It>>>
	void sort(It begin, It end, Comp comp = {})
	{
		return intro_sort(begin, end, comp);
	}

}