#ifndef THREE_PASS_H
#define THREE_PASS_H

/* How the 32 bits gets separated? */
#define TPB1 11 // top
#define TPB2 11 // mid
#define TPB3 10 // bottom

static inline constexpr uint32_t min3u32(uint32_t a, uint32_t b, uint32_t c) {
	return (a <= b) ?
		((a <= c) ? a : c) :
		((b <= c) ? b : c);
}

/** Simple three-pass (ok: 3 + 1) bottom-up radix sort for uint32_t */
static inline void threepass(uint32_t *a, int n) noexcept {
	constexpr int shr1 = TPB3 + TPB2;
	constexpr int shr2 = TPB3;
	constexpr int shr3 = 0;
	constexpr int mask1 = (1 << TPB1) - 1;
	constexpr int mask2 = (1 << TPB2) - 1;
	constexpr int mask3 = (1 << TPB3) - 1;

	/* helper buffers. */
	int sz = n * sizeof(a[0]);

	static thread_local uint32_t bucket1[1 << TPB1];
	memset(bucket1, 0, (1 << TPB1) * sizeof(uint32_t));
	static thread_local uint32_t bucket2[1 << TPB2];
	memset(bucket2, 0, (1 << TPB2) * sizeof(uint32_t));
	static thread_local uint32_t bucket3[1 << TPB3];
	memset(bucket3, 0, (1 << TPB3) * sizeof(uint32_t));

	uint32_t *buf = (uint32_t *)malloc(sz);
	assert(buf != NULL);
	memset(buf, 0, n * sizeof(uint32_t)); // XXX: TODO: REMOVE

	/* Count occurences (can count together with good ILP) */
	#pragma GCC unroll 64
	for(uint32_t i = 0; i < n; ++i) {
		++bucket1[(a[i] >> shr1) & mask1];
		++bucket2[(a[i] >> shr2) & mask2];
		++bucket3[(a[i] >> shr3) & mask3];
	}

	/* Count prefix sums - try as much ILP as possible because bigger arrays than usual! */
	uint32_t prev1 = 0;
	uint32_t prev2 = 0;
	uint32_t prev3 = 0;
	uint32_t common = min3u32(
		(1 << TPB1),
		(1 << TPB2),
		(1 << TPB3)
	);
	int i = 0;
	#pragma GCC unroll 8
	for (; i < common; ++i) {
		bucket1[i] += prev1;
		prev1 = bucket1[i];
		bucket2[i] += prev2;
		prev2 = bucket2[i];
		bucket3[i] += prev3;
		prev3 = bucket3[i];
	}
	/* Do remaining 1 */
	for (int j = i; j < (1 << TPB1); ++j) {
		bucket1[j] += prev1;
		prev1 = bucket1[j];
	}
	/* Do remaining 2 */
	for (int j = i; j< (1 << TPB2); ++j) {
		bucket2[j] += prev2;
		prev2 = bucket2[j];
	}
	/* Do remaining 3 */
	for (int j = i; j < (1 << TPB3); ++j) {
		bucket3[j] += prev3;
		prev3 = bucket3[j];
	}
	
	// Bottom digit
	// right-to-left to ensure already sorted digits order we keep for iterations
	#pragma GCC unroll 64
	for(uint32_t i = n; i > 0; --i) {
		// Prefetch caches
		// Get num and its new offset / location
		auto num = a[i - 1];
		auto bkeyni = (num >> shr3) & mask3;
		auto offset = --bucket3[bkeyni];

		// Add to the proper target location
		buf[offset] = num;
	}
	// Mid digit
	// right-to-left to ensure already sorted digits order we keep for iterations
	#pragma GCC unroll 64
	for(uint32_t i = n; i > 0; --i) {
		// Prefetch caches
		// Get num and its new offset / location
		auto num = buf[i - 1];
		auto bkeyni = (num >> shr2) & mask2;
		auto offset = --bucket2[bkeyni];

		// Add to the proper target location
		a[offset] = num;
	}
	// Top digit
	// right-to-left to ensure already sorted digits order we keep for iterations
	#pragma GCC unroll 64
	for(uint32_t i = n; i > 0; --i) {
		// Prefetch caches
		// Get num and its new offset / location
		auto num = a[i - 1];
		auto bkeyni = (num >> shr1) & mask1;
		auto offset = --bucket1[bkeyni];

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

	// Memcpy back!
	memcpy(a, buf, n * sizeof(uint32_t));

	free(buf);
}

#endif /* THREE_PASS_H */