#include "morton_code.h"

#include <gtest/gtest.h>
#include <stdio.h>
#include <math.h>
#include <random>

#if defined(__GNUC__) && defined(__GNUC_MINOR__) && (__GNUC__ == 4) && (__GNUC_MINOR__ <= 4)
	#define default_random_engine mt19937
	#define uniform_int_distribution uniform_int
#endif

class MortonCodeInsider : public MortonCode
{
public:
	MortonCodeInsider()
	: MortonCode() {}

	MortonCodeInsider(unsigned int xIndex, unsigned int yIndex, unsigned int zIndex, unsigned char lvl)
	: MortonCode(xIndex, yIndex, zIndex, lvl) {}

	MortonCodeInsider(double x, double y, double z, unsigned char lvl)
	: MortonCode(x, y, z, lvl) {}

	unsigned int getKey(int index)
	{
		return key[index];
	}
};

unsigned int maxIndex(unsigned char level)
{
	assert(level <= MORTON_CODE_MAX_LEVEL);
	return (unsigned int) (((unsigned long long) 1 << level) - 1);
}

double getWidth(unsigned char level)
{
	return maxIndex(level) + 1.0;
}

TEST(morton_code, corner_fixed_cases)
{
	{
		unsigned int level = MORTON_CODE_MAX_LEVEL;
		unsigned int x = maxIndex(level);
		unsigned int y = maxIndex(level);
		unsigned int z = maxIndex(level);
		double cube_width = getWidth(level);
		vector3d xyz(x, y, z);
		vector3d center = (xyz + vector3d(0.5, 0.5, 0.5)) / cube_width;

		{
			MortonCodeInsider code(x, y, z, level);
			EXPECT_EQ(code.getKey(0), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(1), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(2), 0xFFFFFFFF);
			EXPECT_EQ(code.center(), center);
		}

		{
			MortonCodeInsider code(center.x(), center.y(), center.z(), level);
			EXPECT_EQ(code.getKey(0), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(1), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(2), 0xFFFFFFFF);
			EXPECT_EQ(code.center(), center);
		}
	}
	//__________________________________________________________________________________________________________________
	{
		unsigned int level = MORTON_CODE_MAX_LEVEL;
		unsigned int x = maxIndex(level) - 1; // ...110
		unsigned int y = maxIndex(level) - 2; // ...101   =>   .........110101011
		unsigned int z = maxIndex(level) - 4; // ...011
		double cube_width = getWidth(level);
		vector3d xyz(x, y, z);
		vector3d center = (xyz + vector3d(0.5, 0.5, 0.5)) / cube_width;

		{
			MortonCodeInsider code(x, y, z, level);
			EXPECT_EQ(code.getKey(0), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(1), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(2), 0xFFFFFFAB); // .........110101011
			EXPECT_EQ(code.center(), center);
		}

		{
			MortonCodeInsider code(center.x(), center.y(), center.z(), level);
			EXPECT_EQ(code.getKey(0), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(1), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(2), 0xFFFFFFAB); // .........110101011
			EXPECT_EQ(code.center(), center);
		}
	}
	//__________________________________________________________________________________________________________________
	{
		unsigned int level = MORTON_CODE_MAX_LEVEL;
		unsigned int x = maxIndex(level) - 4; // 11...011
		unsigned int y = maxIndex(level) - 1; // 11...110   =>   111111.........011110101
		unsigned int z = maxIndex(level) - 2; // 11...101
		double cube_width = getWidth(level);
		vector3d xyz(x, y, z);
		vector3d center = (xyz + vector3d(0.5, 0.5, 0.5)) / cube_width;

		{
			MortonCodeInsider code(x, y, z, level);
			EXPECT_EQ(code.getKey(0), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(1), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(2), 0xFFFFFEF5); // 111111.........011110101
			EXPECT_EQ(code.center(), center);
		}

		{
			MortonCodeInsider code(center.x(), center.y(), center.z(), level);
			EXPECT_EQ(code.getKey(0), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(1), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(2), 0xFFFFFEF5); // 111111.........011110101
			EXPECT_EQ(code.center(), center);
		}
	}
	//__________________________________________________________________________________________________________________
	{
		unsigned int level = MORTON_CODE_MAX_LEVEL;
		unsigned int x = maxIndex(level) - 4 - (1 << 31); // 01...011
		unsigned int y = maxIndex(level) - 1;             // 11...110   =>   011111.........011110101
		unsigned int z = maxIndex(level) - 2;             // 11...101
		double cube_width = getWidth(level);
		vector3d xyz(x, y, z);
		vector3d center = (xyz + vector3d(0.5, 0.5, 0.5)) / cube_width;

		{
			MortonCodeInsider code(x, y, z, level);
			EXPECT_EQ(code.getKey(0), 0x7FFFFFFF);
			EXPECT_EQ(code.getKey(1), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(2), 0xFFFFFEF5); // 011111.........011110101
			EXPECT_EQ(code.center(), center);
		}

		{
			MortonCodeInsider code(center.x(), center.y(), center.z(), level);
			EXPECT_EQ(code.getKey(0), 0x7FFFFFFF);
			EXPECT_EQ(code.getKey(1), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(2), 0xFFFFFEF5); // 011111.........011110101
			EXPECT_EQ(code.center(), center);
		}
	}
	//__________________________________________________________________________________________________________________
	{
		unsigned int level = MORTON_CODE_MAX_LEVEL;
		unsigned int x = maxIndex(level) - 4;             // 11...011
		unsigned int y = maxIndex(level) - 1 - (1 << 31); // 01...110   =>   101111.........011110101
		unsigned int z = maxIndex(level) - 2;             // 11...101
		double cube_width = getWidth(level);
		vector3d xyz(x, y, z);
		vector3d center = (xyz + vector3d(0.5, 0.5, 0.5)) / cube_width;

		{
			MortonCodeInsider code(x, y, z, level);
			EXPECT_EQ(code.getKey(0), 0xBFFFFFFF);
			EXPECT_EQ(code.getKey(1), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(2), 0xFFFFFEF5); // 101111.........011110101
			EXPECT_EQ(code.center(), center);
		}

		{
			MortonCodeInsider code(center.x(), center.y(), center.z(), level);
			EXPECT_EQ(code.getKey(0), 0xBFFFFFFF);
			EXPECT_EQ(code.getKey(1), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(2), 0xFFFFFEF5); // 101111.........011110101
			EXPECT_EQ(code.center(), center);
		}
	}
	//__________________________________________________________________________________________________________________
	{
		unsigned int level = MORTON_CODE_MAX_LEVEL;
		unsigned int x = maxIndex(level) - 4;             // 11...011
		unsigned int y = maxIndex(level) - 1;             // 11...110   =>   110111.........011110101
		unsigned int z = maxIndex(level) - 2 - (1 << 31); // 01...101
		double cube_width = getWidth(level);
		vector3d xyz(x, y, z);
		vector3d center = (xyz + vector3d(0.5, 0.5, 0.5)) / cube_width;

		{
			MortonCodeInsider code(x, y, z, level);
			EXPECT_EQ(code.getKey(0), 0xDFFFFFFF);
			EXPECT_EQ(code.getKey(1), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(2), 0xFFFFFEF5); // 110111.........011110101
			EXPECT_EQ(code.center(), center);
		}

		{
			MortonCodeInsider code(center.x(), center.y(), center.z(), level);
			EXPECT_EQ(code.getKey(0), 0xDFFFFFFF);
			EXPECT_EQ(code.getKey(1), 0xFFFFFFFF);
			EXPECT_EQ(code.getKey(2), 0xFFFFFEF5); // 110111.........011110101
			EXPECT_EQ(code.center(), center);
		}
	}
}

void testCube(unsigned int x, unsigned int y, unsigned int z, unsigned char level)
{
	double cube_width = getWidth(level);
	vector3d xyz(x, y, z);
	vector3d center = (xyz + vector3d(0.5, 0.5, 0.5)) / cube_width;

	MortonCode code(x, y, z, level);
	EXPECT_EQ(code.center(), center);
	{
		MortonCode code2(center.x(), center.y(), center.z(), level);
		EXPECT_EQ(code, code2);
	}

	if (level > 0) {
		MortonCode parent0(x / 2, y / 2, z / 2, level - 1);
		MortonCode parent1 = code.createParent();
		EXPECT_EQ(parent0, parent1);

		for (unsigned int dx = 0; dx < 2; ++dx) {
			for (unsigned int dy = 0; dy < 2; ++dy) {
				for (unsigned int dz = 0; dz < 2; ++dz) {
					unsigned int x0 = x / 2 * 2;
					unsigned int y0 = y / 2 * 2;
					unsigned int z0 = z / 2 * 2;
					MortonCode sibling0(x0 + dx, y0 + dy, z0 + dz, level);
					MortonCode sibling1 = code.createSibling(dx, dy, dz);
					MortonCode sibling2 = parent1.createChild(dx, dy, dz);
					EXPECT_EQ(sibling0, sibling1);
					EXPECT_EQ(sibling1, sibling2);

					unsigned int dx0 = std::numeric_limits<unsigned int>::max();
					unsigned int dy0 = std::numeric_limits<unsigned int>::max();
					unsigned int dz0 = std::numeric_limits<unsigned int>::max();
					sibling0.getSubindexFromParent(dx0, dy0, dz0);
					EXPECT_EQ(dx0, dx);
					EXPECT_EQ(dy0, dy);
					EXPECT_EQ(dz0, dz);
				}
			}
		}

		for (int dz = 0; dz <= 2; ++dz) {
			for (int dy = 0; dy <= 2; ++dy) {
				for (int dx = 0; dx <= 2; ++dx) {
					bool neighbour_out_of_cube;
					MortonCode neighbour1 = code.createNeighbour(dx, dy, dz, neighbour_out_of_cube);
					EXPECT_EQ(neighbour_out_of_cube,
							  (ptrdiff_t) x + dx - 1 < 0 || (ptrdiff_t) x + dx - 1 > maxIndex(level) ||
							  (ptrdiff_t) y + dy - 1 < 0 || (ptrdiff_t) y + dy - 1 > maxIndex(level) ||
							  (ptrdiff_t) z + dz - 1 < 0 || (ptrdiff_t) z + dz - 1 > maxIndex(level));
					if (neighbour_out_of_cube)
						continue;

					vector3d expected_center = (vector3d(x + dx - 1, y + dy - 1, z + dz - 1) + vector3d(0.5, 0.5, 0.5)) / cube_width;
					vector3d found_center = neighbour1.center();
					EXPECT_EQ(found_center, expected_center);
//					double voxel_radius = 0.5 / cube_width;
//					EXPECT_LT((found_center - expected_center).norm(), 0.1 * voxel_radius);

					MortonCode neighbour0(x + dx - 1, y + dy - 1, z + dz - 1, level);
					EXPECT_EQ(neighbour0, neighbour1);
				}
			}
		}
	}
}

TEST(morton_code, fixed_cases)
{
	unsigned int level = 23;
	unsigned int x = 262144; // 0100 0000 0000 0000 0000
	unsigned int y = 131072; // 0010 0000 0000 0000 0000
	unsigned int z = 65536;  // 0001 0000 0000 0000 0000

	testCube(x, y, z, level);
}

TEST(morton_code, simple_cases)
{
	for (int level = 0; level < 8; ++level) {
		for (unsigned int x = 0; x <= maxIndex(level); ++x) {
			for (unsigned int y = 0; y <= maxIndex(level); ++y) {
				for (unsigned int z = 0; z <= maxIndex(level); ++z) {
					testCube(x, y, z, level);
				}
			}
		}
	}
	for (int level = 8; level <= MORTON_CODE_MAX_LEVEL; ++level) {
		unsigned long long step = maxIndex(level);
		step = (step + 7) / 8;
		for (unsigned int xi = 0; xi <= 8; ++xi) {
			for (unsigned int yi = 0; yi <= 8; ++yi) {
				for (unsigned int zi = 0; zi <= 8; ++zi) {
					unsigned int x = std::min(xi * step, (unsigned long long) maxIndex(level));
					unsigned int y = std::min(yi * step, (unsigned long long) maxIndex(level));
					unsigned int z = std::min(zi * step, (unsigned long long) maxIndex(level));
					testCube(x, y, z, level);
				}
			}
		}
	}
}

TEST(morton_code, random_cases)
{
	int n = 2390000;
	std::default_random_engine r(239);
	std::uniform_int_distribution<int> level_distribution(9, MORTON_CODE_MAX_LEVEL);

	for (int i = 0; i < n; ++i) {
		int level = level_distribution(r);

		std::uniform_int_distribution<unsigned int> xyz_distribution(0, maxIndex(level));
		unsigned int x = xyz_distribution(r);
		unsigned int y = xyz_distribution(r);
		unsigned int z = xyz_distribution(r);

		testCube(x, y, z, level);
	}
}