#ifndef MAGYAR_SORT_H
#define MAGYAR_SORT_H
/**
 * single header lib: In-place, fast heavily modified and optimized radix sort.
 *
 * Only unsigned ints for now, but should be able to modify for int and float...
 * This is the counting variant with smart changes (not per-bit).
 *
 * LICENCE: CC3 - look it up, you need to mention me but that is all
 */

#include <cstdio>
#include <cstdint>
#include <cstring> // memset

// TODO: Only for the regular radix I guess
#include <vector>
#include <algorithm> // std::swap

namespace MagyarSort {
	/* CONFIG */

	// Only change these if you know what you are doing
	// I use these because I want to see if nibbles are
	// better or something...
	//
	// Bytes of nibbles only:
	// - DIGIT_RANGE and BITS_PER_DIGIT should correspond
	// - DIGITS should also correspond with the uint32_t
	// - and DIGIT_RANGE should be 2^n value (16 or 256)
#ifdef MAGYAR_SORT_NIBBLE
	// Per-nibble digits sorting
	static constexpr int DIGITS = 8; // "helyiérték"
	static constexpr int BITS_PER_DIGIT = 4; // "bit / helyiérték"
	static constexpr int DIGIT_RANGE = 16; // "helyiérték állapottér"
#else
	// Per-byte digits sorting
	static constexpr int DIGITS = 4; // "helyiérték"
	static constexpr int BITS_PER_DIGIT = 8; // "bit / helyiérték"
	static constexpr int DIGIT_RANGE = 256; // "helyiérték állapottér"
#endif

	/* DEBUG */

	void debugArr(uint32_t *arr, size_t size) {
		for(int i = 0; i < size; ++i) {
			printf("%x, ", arr[i]);
		}
		printf("\n");
	}

	void debugRadics(size_t *radics) {
		for(size_t j = 0; j < DIGITS; ++j) {
			printf("d%zu: ", j);
			for(size_t i = 0; i < DIGIT_RANGE; ++i) {
				printf("%zu,", radics[i + DIGIT_RANGE*j]);
			}
			printf("\n\n");
		}
	}

	/* HELPERS */

	template<int DIGIT_CHOICE>
	static inline __attribute__((always_inline)) uint32_t getDigit(uint32_t num) noexcept {
		static constexpr int SHIFT = DIGIT_CHOICE * BITS_PER_DIGIT;

		uint32_t shifted = num >> SHIFT;
		return shifted & (DIGIT_RANGE - 1);
	}

	/** Recursive Functor: no class should be generated I think (compiler should be smart) */
	template<int DIGIT>
	struct OccurenceMagic : public OccurenceMagic<DIGIT - 1> {
		inline __attribute__((always_inline)) OccurenceMagic(uint32_t arr[], size_t i, size_t *radicsOut) noexcept
				: OccurenceMagic<DIGIT -1 >(arr, i, radicsOut) {
			// Parents run first so template recursion runs DIGIT=0 first...
			++radicsOut[getDigit<DIGIT>(arr[i]) + DIGIT_RANGE * DIGIT];
		}
	};
	/** Ends template recursion */
	template<>
	struct OccurenceMagic<-1> {
		inline OccurenceMagic(uint32_t arr[], size_t i, size_t *radicsOut) noexcept {}
	};

	static inline void countOccurences(uint32_t arr[], size_t size, size_t *radicsOut) noexcept {
		#pragma GCC unroll 64
		for(size_t i = 0; i < size; ++i) {
			// Creates no object, struct is empty
			OccurenceMagic<DIGITS - 1>(arr, i, radicsOut);
		}
	}

	/** Recursive Functor: no class should be generated I think (compiler should be smart) */
	template<int DIGIT>
	struct PrefixMagic : public PrefixMagic<DIGIT - 1> {
		inline __attribute__((always_inline)) PrefixMagic(size_t *radics, size_t *prev, int i) noexcept
				: PrefixMagic<DIGIT - 1>(radics, prev, i) {
			static constexpr int DSTART = (DIGIT * DIGIT_RANGE);
			radics[DSTART + i] += prev[DIGIT];
			prev[DIGIT] = radics[DSTART + i];
		}
	};
	/** Ends template recursion */
	template<>
	struct PrefixMagic<-1> {
		inline PrefixMagic(size_t *radics, size_t *prev, int i) noexcept {}
	};

	/** Gets REFERENCE to the given digit from the radix-array that has more than one digits */
	template<int DIGIT>
	static inline __attribute__((always_inline)) size_t &rGet(size_t *radics, size_t i) noexcept {
		static constexpr int DSTART = (DIGIT * DIGIT_RANGE);
		return radics[DSTART + i];
	}

	static inline void calcPrefixSums(size_t *radics) noexcept {
		static thread_local size_t prev[DIGITS];
		memset(prev, 0, sizeof(prev));

		for(int i = 0; i < DIGIT_RANGE; ++i) {
			// This is a template-unrolled loop too
			PrefixMagic<DIGITS - 1>(radics, prev, i);
		}
	}

	/** Recursive Functor: no class should be generated I think (compiler should be smart) */
	template<int DIGIT>
	struct RadixMagic : public RadixMagic<DIGIT - 1> {
		inline __attribute__((always_inline)) RadixMagic(size_t *radics, uint32_t *&from, uint32_t *&to, size_t size) noexcept // BEWARE: "*&" needed to swap pointers..
				: RadixMagic<DIGIT - 1>(radics, from, to, size) {
			// DEBUG
			//printf("%d before: ", DIGIT);
			//debugArr(from, size);

			#pragma GCC unroll 64
			for(size_t i = size; i > 0; --i) { // right-to-left to ensure already sorted digits order we keep for iterations
				// Get num and its new offset / location
				auto num = from[i - 1];
				auto digVal = getDigit<DIGIT>(num);
				auto offset = (--rGet<DIGIT>(radics, digVal));

				// Add to the proper target location
				to[offset] = num;
			}

			// DEBUG
			//printf("%d after: ", DIGIT);
			//debugArr(to, size);

			// Only swaps pointers :-)
			std::swap(from, to);
		}
	};
	/** Ends template recursion */
	template<>
	struct RadixMagic<-1> {
		inline RadixMagic(size_t *radics, uint32_t *&from, uint32_t *&to, size_t size) noexcept { }
	};

	/* SORT */

	/**
	 * Example: A simple "vector-giver" which provides a static thread_local that is reused
	 *
	 * This is to be used when you will call sort many times successively!
	 * If you forget to garbage-collect manually, use a VectorGiverHeap.
	 *
	 * XXX - BEWARE: This give references - that is also acceptable and supported!
	 *
	 * This is thread-safe (the Heap one also).
	 */
	struct VectorGiverWithReuse {
		/**
		 * Give a reference to the vector to use as temporary.
		 * Will be resized, is reused so "leaks" memory to be the biggest sorted array size, but you can "Gc()".
		 *
		 * @param s The given vector should have this size.
		 * @param gc OPTIONAL: When true, we create a new empty shared vector. This saves memory after a big sort!
		 * @returns A reference that never go out of scope!
		 */
		static inline __attribute__((always_inline)) std::vector<uint32_t> &Give(size_t s, const bool gc = false) noexcept {
			static thread_local std::vector<uint32_t> arc(s); // saves time on first call to have size here!
			if(gc) { arc = std::vector<uint32_t>(); } // by default optimized out!
			arc.resize(s); // JHP
			// Safe because of static it will not go out of scope
			return arc; // just a reference - no copy!
		}

		/** Release memory back to zero. After this, the first sort will need memory from heap again. */
		inline __attribute__((always_inline)) void Gc() noexcept {
			VectorGiverWithReuse::Give(0, true);
		}
	};

	/**
	 * Example: A simple "vector-giver" which provides new vector from heap.
	 *
	 * This is thread-safe (the VectorGiverWithReuse one also).
	 */
	struct VectorGiverHeap {
		/**
		 * Give a temporary vector which is to be created on heap and freed after sort.
		 *
		 * XXX - BEWARE: Please mind we do not return reference, but value here!
		 *               This works because standard ENSURES return value optimization!
		 *
		 * @param s The given vector should have this size.
		 * @param gc OPTIONAL: When true, we create a new empty shared vector. This saves memory after a big sort!
		 * @returns A vector of appropriate size.
		 */
		inline __attribute__((always_inline)) std::vector<uint32_t> Give(size_t s) noexcept {
			return std::vector<uint32_t>(s); // RVO ensured!
		}

	};

	/*
	 * Sort the given array (in-place sorting) with the given size.
	 *
	 * Rem.: If you use the VectorGiverWithReuse please remind yourself to Gc() it time-to-time!
	 *
	 * Beware: GC needs to happen on all threads that use us!
	 *
	 * @param arr The array to sort. Result will be in the same array - as sorted.
	 * @param size The lenght of the array.
	 * @param REUSE OPTIONAL: When true, we reuse the array instead of always gettin' and releasin' from da heap.
	 * @param GC OPTIONAL: When true, we garbage collect memory from previous sorts if REUSE is true.
	 * @param GC_WITHOUT_SORT OPTIONAL: When true, we "just GC" but do not sort in case of GC is true.
	 */
	template<bool REUSE = false, bool GC = false, bool GC_WITHOUT_SORT = false>
	inline void __attribute__((always_inline)) sort_impl(uint32_t arr[], size_t size) noexcept {
		// Most funny optimization is this multiply here :-)
		//
		// Literally.. come on.. this makes it nearly a compile-time, macro-like
		// ifdef-like thing as we avoid memory allocations of size BUT also we
		// optimize the first call for sort when we REUSE the array so size is fine!
		static thread_local std::vector<uint32_t> arc(size * REUSE);

		// "Garbage-collection"
		if(GC) {
			arc = std::vector<uint32_t>();
			// This must be implemented, because we can only access
			// the static in our function body so this is the "way".
			if(GC_WITHOUT_SORT) {
				return;
			}
		}

		// Holds "digit" occurences, prefix sums, whatevers
		// First "DIGIT_RANGE" elem is for MSB "DIGITS", last is for LSB
		static thread_local size_t radics[DIGITS * DIGIT_RANGE];
		memset(radics, 0, sizeof(radics));

		// Calculate occurences of digits
		countOccurences(arr, size, radics);

		//debugRadics(radics);

		// Calculate prefix sums
		calcPrefixSums(radics);

		//debugRadics(radics);

		/* Regular (old) radix sort with small twist */

		// Regular radix sort - I just changed occurence couting and prefix summing to have more ILP
		// But because my approach does not use that, I want to keep this version in a branch for a
		// regular radix sort using better ILP just to see how it is doing if I wrote those "Magic"
		// above already anyways...

		// Regular radix sort needs a copy, see: https://www.youtube.com/watch?v=ujb2CIWE8zY
		// But instead of the below, we do a trickery...
		//
		//std::vector<uint32_t> arc(size);
		//auto arc = VectorGiver::Give(size); // "auto" is needed for this to perform well with some givers!
		//
		// Rem.: The branch is optimized out in compile time!
		if(REUSE) {
			arc.resize(size);
		} else {
			// Must not be .clean() !!!
			// We must regain memory of previous!
			arc = std::move(std::vector<uint32_t>(size));
		}

		uint32_t *from = arr;
		uint32_t *to = &arc[0];

		RadixMagic<DIGITS - 1>(radics, from, to, size);

		// With an other API we could spare this copy if we can delete original arr and return ptr or something...
		// I am fine with this... this is not my main idea anyways, just little ILP tweak to regular radix sort
		//if(to != arr) // <- logically, but bad they are already swapped here!!! BEWARE
		if(from != arr) { // <- in reality this is what we want because of last swap happened anyways!
			memcpy(arr, from, size);
		}
	}

	/**
	 * Garbage collect reused data structures from last call.
	 *
	 * This is optimized and is a NO-OP if MAGYAR_SORT_DEFAULT_REUSE is not defined!
	 *  - unless you use the FORCE! May it be with you if you need it.
	 *
	 * @param FORCE OPTIONAL: When true, the gc happens even if MAGYAR_SORT_DEFAULT_REUSE is not defined!
	 */
	template<bool FORCE = false>
	inline void gc() noexcept {
		if(FORCE) {
			// Only GC-ing
			MagyarSort::sort_impl<true, true, true>(nullptr, 0);
		} else {
#ifdef MAGYAR_SORT_DEFAULT_REUSE
			// Only GC-ing
			MagyarSort::sort_impl<true, true, true>(nullptr, 0);
#endif
		}
	}

	/**
	 * Sort the given array (in-place sorting) with the given size.
	 *
	 * Rem.: Please remind yourself to cc() from time-to-time!
	 * Rem.: Thread-safe to use!
	 *
	 * Beware: MagyarSort::gc<true>(); needs to happen on all threads that use this variant otherwise memory leaks away!
	 *         Please mind the "true" template parameter that forces the GC even when sort by default not reuses...
	 *
	 * @param arr The array to sort. Result will be in the same array - as sorted.
	 * @param size The lenght of the array.
	 * @param GC OPTIONAL: When true, we garbage collect before this sort - so cached memory size will be "size" elems.
	 */
	template<bool GC = false>
	inline void __attribute__((always_inline)) sort_reuse(uint32_t arr[], size_t size) noexcept {
		// Reuse the temporary vectors across runs
		// This results in much less heap allocations and much faster on gcc
		// and also a bit faster on clang too.
		MagyarSort::sort_impl<true>(arr, size);
	}

	/**
	 * Sort the given array (in-place sorting) with the given size.
	 *
	 * Rem.: Thread-safe to use!
	 *
	 * Beware: MagyarSort::gc(); needs to happen on all threads that use this variant otherwise memory leaks away!
	 *
	 * @param arr The array to sort. Result will be in the same array - as sorted.
	 * @param size The lenght of the array.
	 */
	inline void __attribute__((always_inline)) sort_no_reuse(uint32_t arr[], size_t size) noexcept {
		// We use the heap once per every call...
		// This is safer and we do not need garbage collecting
		MagyarSort::sort_impl(arr, size);
	}

	/*
	 * Sort the given array (in-place sorting) with the given size.
	 *
	 * Rem.: If you use the VectorGiverWithReuse please remind yourself to Gc() it time-to-time!
	 *
	 * Beware: MagyarSort::gc(); should be called after "sort bursts" (consecutive fast sorts of when you need memory
	 *         on all threads that use this variant otherwise memory leaks away as biggest sorted array keeps being in ram!
	 *         This depends on the config #define MAGYAR_SORT_DEFAULT_REUSE is defined or not. Define and you get reuse
	 *         and if you get reuse you can call multiple sorts with reused temporary buffers that you gc() afterwards!
	 *
	 * @param arr The array to sort. Result will be in the same array - as sorted.
	 * @param size The lenght of the array.
	 */
	inline void sort(uint32_t arr[], size_t size) noexcept {
#ifdef MAGYAR_SORT_DEFAULT_REUSE
		MagyarSort::sort_reuse(arr, size);
#else
		MagyarSort::sort_no_reuse(arr, size);
#endif
	}
};

#endif