#include <cstdio>
#include <cassert>
#include <vector>
#include <string>
#include <chrono>
#include "amap.h"
#include "simap.h"
#include "mapmap.hpp"
#include "unomap.hpp"
#include "simd_map.h"
#include "vmap.h"

/**
 * Creates keys or returns the ith key. Used for performance tests.
 *
 * @param i When "create" is false, we return the ith key (does not check OOB)
 * @param create When true, we initialize the keystore with keys generated from 0..i indices.
 * @returns The ith key when create==false, otherwise undefined.
 */
inline const char *keystore(int i, bool create = false) noexcept {
	static thread_local std::vector<std::string> keys;

	if(!create) {
		return keys[i].c_str();
	} else {
		keys.resize(0);
		keys.reserve(0);
		std::string key = "k";
		for(int j = 0; j < i; ++j) {
			keys.push_back(key + std::to_string(j));
		}
		return NULL;
	}
}

/**
 * Creates keys or returns the ith integer key. Used for performance tests.
 *
 * Rem.: Generated keys are like this: i, i-1, i-2, ... 1
 *
 * @param i When "create" is false, we return the ith key (does not check OOB)
 * @param create When true, we initialize the keystore with keys generated from 0..i indices.
 * @returns The ith key when create==false, otherwise undefined.
 */
inline int *int_keystore(int i, bool create = false) noexcept {
	static thread_local std::vector<int> keys;

	if(!create) {
		return &(keys[i]);
	} else {
		keys.resize(0);
		keys.reserve(0);
		for(int j = 0; j < i; ++j) {
			keys.push_back(i - j);
		}
		return NULL;
	}
}
inline const char *datastore_int(int i, bool create = false) noexcept {
	static thread_local std::vector<std::string> keys;

	if(!create) {
		return keys[i].c_str();
	} else {
		keys.resize(0);
		keys.reserve(0);
		std::string key = "k";
		for(int j = 0; j < i; ++j) {
			keys.push_back(key + std::to_string(j));
		}
		return NULL;
	}
}

/**
 * Creates datas or returns the ith data. Used for performance tests.
 *
 * @param i When "create" is false, we return the ith data (does not check OOB)
 * @param create When true, we initialize the datastore with datas generated from 0..i indices.
 * @returns The ith data when create==false, otherwise undefined.
 */
inline int *datastore(int i, bool create = false) noexcept {
	static thread_local std::vector<int> keys;

	if(!create) {
		return &(keys[i]);
	} else {
		keys.resize(0);
		keys.reserve(0);
		for(int j = 0; j < i; ++j) {
			keys.push_back(j);
		}
		return NULL;
	}
}

void test_perf(amap mapdo, void *map, int max_key, const char *what) {
	auto begin = std::chrono::high_resolution_clock::now();
	for(int i = 0; i < max_key; ++i) {
		const char *key = keystore(i);
		int *data = datastore(i);
		mapdo(map, AMAP_SET, key, data);
	}
	auto end = std::chrono::high_resolution_clock::now();
	auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);

	printf("Insertion time for %d elements (%s): %.3f ms.\n", max_key, what, elapsed.count() * 1e-6);
}

void test_basics(amap mapdo, void *map) {
	/* Most basics */
	assert(NULL == mapdo(map, AMAP_GET, "asdf", NULL));
	int i = 42;
	int *iptr;
	const char *chptr;
	assert(NULL != mapdo(map, AMAP_SET, "meaning", &i));
	assert(NULL != (iptr = (int *)mapdo(map, AMAP_GET, "meaning", NULL)));
	assert(*iptr == 42);
	assert(iptr == &i);

	/* Delete / tombstone */
	assert(NULL != mapdo(map, AMAP_SET, "meaning", NULL));
	assert(NULL == (int *)mapdo(map, AMAP_GET, "meaning", NULL));

	/* Check re-adding */
	assert(NULL != mapdo(map, AMAP_SET, "meaning", &i));
	assert(NULL != (iptr = (int *)mapdo(map, AMAP_GET, "meaning", NULL)));
	assert(*iptr == 42);
	assert(iptr == &i);

	/* Test Erase */
	assert(NULL != mapdo(map, AMAP_ERASE, NULL, NULL));

	/* Check re-adding 3 new things */
	assert(NULL != mapdo(map, AMAP_SET, "meaningless1", &i));
	assert(NULL != mapdo(map, AMAP_SET, "meaning2", &i));
	const char *helloworld = "Hello world!";
	assert(NULL != mapdo(map, AMAP_SET, "hello", (char *)helloworld)); /* ugly cast... */

	assert(NULL != (chptr = (const char *)mapdo(map, AMAP_GET, "hello", NULL)));
	assert(strlen(chptr) == strlen(helloworld));
	assert(strcmp(chptr, helloworld) == 0);
	assert(NULL != (iptr = (int *)mapdo(map, AMAP_GET, "meaning2", NULL)));
	assert(*iptr == 42);
	assert(iptr == &i);
	assert(NULL != (iptr = (int *)mapdo(map, AMAP_GET, "meaningless1", NULL)));
	assert(*iptr == 42);
	assert(iptr == &i);

	/* Check the case where we have same 8-long prefix for multiple and they should be different */
	int long_1 = 1;
	int long_2 = 2;
	int long_3 = 3;

	assert(NULL != mapdo(map, AMAP_SET, "very_long_test_key_1", &long_1));
	assert(NULL != mapdo(map, AMAP_SET, "very_long_test_key_2", &long_2));
	assert(NULL != mapdo(map, AMAP_SET, "very_long_test_key_3", &long_3));

	assert(NULL != (iptr = (int *)mapdo(map, AMAP_GET, "very_long_test_key_1", NULL)));
	assert(*iptr == 1);
	assert(NULL != (iptr = (int *)mapdo(map, AMAP_GET, "very_long_test_key_2", NULL)));
	assert(*iptr == 2);
	assert(NULL != (iptr = (int *)mapdo(map, AMAP_GET, "very_long_test_key_3", NULL)));
	assert(*iptr == 3);
}

void test_stringmaps(int perf_test_i) {
	/* Basic tests */
	simap_instance si = simap_create();
	test_basics(simap, &si);

	mapmap_instance mi = mapmap_create();
	test_basics(mapmap, &mi);

	unomap_instance umi = unomap_create();
	test_basics(unomap, &umi);

	/* Performance tests */
	int i = perf_test_i;
	test_perf(mapmap, &mi, i, "std::map");
	test_perf(simap, &si, i, "simap");
	test_perf(unomap, &umi, i, "std::unordered_map");
}

void test_simd_map_basics() {
	/* Empty free tests */
	simd_map smap = simd_map_create();
	simd_map_free(&smap);

	/* Re-create */
	smap = simd_map_create();

	/* Empty search */
	assert(simd_map_find(&smap, 42) == NULL);
	assert(simd_map_size(&smap) == 0);

	/* Insertions */
	assert(simd_map_set(&smap, 40, 0) != 0);
	assert(simd_map_set(&smap, 41, 1) != 0);
	assert(simd_map_set(&smap, 42, 2) != 0);
	assert(simd_map_size(&smap) == 3);

	/* Searches */
	assert(*simd_map_find(&smap, 40) == 0);
	assert(*simd_map_find(&smap, 41) == 1);
	assert(*simd_map_find(&smap, 42) == 2);

	/* Test erase */
	simd_map_erase(&smap);
	assert(simd_map_find(&smap, 42) == NULL);
	assert(simd_map_set(&smap, 42, 2) != 0);
	assert(*simd_map_find(&smap, 42) == 2);

	/* Test a bit more */
	int cnt = 100;
	for(int i = 0; i < cnt; ++i) {
		assert(simd_map_set(&smap, i, (cnt - i)) != 0);
	}
	assert(simd_map_size(&smap) == 100); /* 42->2 should get overwritten */
	for(int i = 0; i < cnt; ++i) {
		assert(*simd_map_find(&smap, i) == (uint32_t)(cnt - i));
	}

	/* Test removal */
	assert(simd_map_remove(&smap, 41) != 0);
	assert(simd_map_remove(&smap, 43) != 0);
	assert(simd_map_find(&smap, 41) == NULL);
	assert(simd_map_find(&smap, 43) == NULL);
	assert(simd_map_find(&smap, 42) != NULL);
	assert(simd_map_find(&smap, 99) != NULL);
	assert(simd_map_find(&smap, 98) != NULL);
	assert(simd_map_size(&smap) == 98);

	/* Test multimap operations */
	assert(simd_map_multi_set(&smap, 42, 42) != 0);
	assert(simd_map_multi_set(&smap, 42, 43) != 0);
	assert(simd_map_find(&smap, 42) != NULL);
	assert(*simd_map_find(&smap, 42) != 42);
	simd_map_find_res res1 = simd_map_find_all(&smap, 42, simd_map_find_all_begin());
	assert(res1.value_location != NULL);
	assert(*(res1.value_location) == (uint32_t)(cnt - 42));
	simd_map_find_res res2 = simd_map_find_all(&smap, 42, res1);
	assert(*(res2.value_location) == 42);
	simd_map_find_res res3 = simd_map_find_all(&smap, 42, res2);
	assert(res3.value_location != NULL);
	assert(*(res3.value_location) == 43);

	/* Test filled-free */
	simd_map_free(&smap);
}

void test_simd_map_perf(int max_key) {
#ifdef __AVX2__
	puts("...Perf testing simd_map with AVX2...");
#elif __SSE2__
	puts("...Perf testing simd_map with SSE2...");
#endif

	// XXX: This way we would measure the wrong thing:
	// simd_map smap = simd_map_create();
	// Why? To measure the right thing, not the first allocation!
	simd_map smap = simd_map_create_and_reserve();
	auto begin = std::chrono::high_resolution_clock::now();
	for(int i = 0; i < max_key; ++i) {
		int *key = int_keystore(i);
		int *data = datastore(i);
		simd_map_set(&smap, *key, *data);
	}
	auto end = std::chrono::high_resolution_clock::now();
	auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);

	printf("Insertion time for %d elements (simd_map): %.3f ms.\n", max_key, elapsed.count() * 1e-6);
	simd_map_free(&smap);
}

void test_int_unordered_map(int max_key) {
	std::unordered_map<int, int> smap;
	auto begin = std::chrono::high_resolution_clock::now();
	for(int i = 0; i < max_key; ++i) {
		int *key = int_keystore(i);
		int *data = datastore(i);
		smap[*key] = *data;
	}
	auto end = std::chrono::high_resolution_clock::now();
	auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);

	printf("Insertion time for %d elements (std::unordered_map<int,int>): %.3f ms.\n", max_key, elapsed.count() * 1e-6);
}

void test_int_std_map(int max_key) {
	std::map<int, int> smap;
	auto begin = std::chrono::high_resolution_clock::now();
	for(int i = 0; i < max_key; ++i) {
		int *key = int_keystore(i);
		int *data = datastore(i);
		smap[*key] = *data;
	}
	auto end = std::chrono::high_resolution_clock::now();
	auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);

	printf("Insertion time for %d elements (std::map<int,int>): %.3f ms.\n", max_key, elapsed.count() * 1e-6);
}

void test_intmaps(int perf_test_i) {
	/* Basic tests */
	// test_simd_map_basics();
	test_int_std_map(perf_test_i);
	test_int_unordered_map(perf_test_i);
	test_simd_map_perf(perf_test_i);
}

int main() {
	int perf_test_i = 1000;

	/* Prepare data stores */
	keystore(perf_test_i, true);
	int_keystore(perf_test_i, true);
	datastore(perf_test_i, true);

	/* Tests */
	puts("");
	puts("Integer maps...");
	puts("");
	test_intmaps(perf_test_i);
	puts("");
	puts("String maps...");
	puts("");
	test_stringmaps(perf_test_i);
	puts("");
	puts("...done!");

	return 0;
}