#include <d2d1_1.h>
#include <vector>
#include <Dwmapi.h>
#include <thread>
#include <memory>
#include <fcntl.h>
#include <io.h>
#include <cstdio>
#include <string>
#include <iostream>
#include <conio.h>
#include <chrono>
#include <algorithm>
#include <random>
#include <map>
#include <utility>
#include <set>
#include <iterator>

#include <D2d1_1.h>
#include <dwrite.h>
#pragma comment(lib, "d2d1")
#pragma comment(lib, "Dwmapi")
#pragma comment(lib, "dwrite")

LRESULT CALLBACK WndProc(HWND, UINT, WPARAM, LPARAM);

HWND hWnd;
ID2D1Factory* factory = NULL;
IDWriteFactory* wfactory = NULL;
ID2D1HwndRenderTarget* target = NULL;
IDWriteTextFormat* format = NULL;

ID2D1Bitmap* m_pBitmap;
RECT WindowRect;

uint32_t board_size = 0;
uint32_t colors = 0;
uint32_t simulations = 0;
uint32_t cell_size = 16;
uint32_t sleep_cycle = 500;
std::vector<ID2D1SolidColorBrush*> brushes;
uint32_t cluster_id = 1;

bool SanitizeInput(std::string input, int lower, int upper, uint32_t& output, uint32_t d_output = 0)
{
	if (input.empty())
	{
		if (d_output)
		{
			printf("Using default value: %d\n", d_output);
			output = d_output;
			return true;
		}
		return false;
	}
	int v = atoi(input.c_str());
	if (v < lower || v > upper)
	{
		printf("Value (%d) is not between %d and %d\n", v, lower, upper);
		return false;
	}
	output = static_cast<uint32_t>(v);
	return true;
}

enum Direction { LEFT, UP, DOWN, RIGHT };

struct Tile
{
	uint32_t color = 0;
	int32_t cluster = 0;	
	std::set<uint32_t> clusters; //neighbors	
	void evaluate()
	{
		std::set<uint32_t>::iterator it = clusters.find(cluster);
		if (it != clusters.end()) clusters.erase(cluster);
	}
};
Tile* board = nullptr;

struct Cluster
{
	//Color, Value
	std::set<uint32_t> clusters; //the clusters it touches
	std::set<uint32_t> node_ids;
	int32_t color;
	void AddTile(uint32_t index)
	{
		const Tile& t = board[index];
		node_ids.insert(index);
	}
	void SetColor(uint32_t color)
	{
		this->color = color;
		for (uint32_t node_id : node_ids)
			board[node_id].color = color;
	}
	void evaluate()
	{
		clusters.clear();
		for (uint32_t id : node_ids)
		{
			Tile& t = board[id];
			clusters.insert(t.clusters.begin(), t.clusters.end());
		}
	}
	Cluster(uint32_t index)
	{
		AddTile(index);
		const Tile& t = board[index];
		color = t.color;
	}

	Cluster() 
	{
		color = 0;
	}
};
std::map<uint32_t, Cluster> cluster_map;

void ClusterTile(uint32_t x, uint32_t y)
{
	Tile& current = board[x + board_size * y];
	std::vector<uint32_t> dfs_stack;
	std::set<uint32_t> dfs_discovered;

	uint32_t cx = x, cy = y;
	if (cx > 0)
	{
		uint32_t next_id = (cx - 1) + board_size * cy;
		if (dfs_discovered.find(next_id) == dfs_discovered.end())
			dfs_stack.push_back(next_id);
	}
	if (cx < board_size - 1)
	{
		uint32_t next_id = (cx + 1) + board_size * cy;
		if (dfs_discovered.find(next_id) == dfs_discovered.end())
			dfs_stack.push_back(next_id);
	}
	if (cy > 0)
	{
		uint32_t next_id = cx + board_size * (cy - 1);
		if (dfs_discovered.find(next_id) == dfs_discovered.end())
			dfs_stack.push_back(next_id);
	}
	if (cy < board_size - 1)
	{
		uint32_t next_id = cx + board_size * (cy + 1);
		if (dfs_discovered.find(next_id) == dfs_discovered.end())
			dfs_stack.push_back(next_id);
	}

	if (current.cluster == 0)
	{
		current.cluster = cluster_id++;
		cluster_map[current.cluster] = Cluster(x + board_size * y);
	}

	std::map<uint32_t, Cluster>::iterator cluster_it;
	cluster_it = cluster_map.find(current.cluster);

	while (!dfs_stack.empty())
	{
		std::vector<uint32_t>::iterator it = dfs_stack.begin();
		uint32_t d_idx = (*it);
		dfs_stack.erase(it);
		dfs_discovered.insert(d_idx);
		uint32_t dx = d_idx % board_size;
		uint32_t dy = d_idx / board_size;

		Tile& compare = board[dx + board_size * dy];
		if (current.color != compare.color)
		{
			std::set<uint32_t>::iterator cluster_it;
			cluster_it = compare.clusters.find(current.cluster);
			if (cluster_it == compare.clusters.end())
				compare.clusters.insert(current.cluster);
			continue;
		}

		if (current.cluster && current.cluster == compare.cluster)
		{
			continue;
		}
		
		compare.cluster = current.cluster;
		compare.evaluate();
		cluster_it->second.AddTile(d_idx);

		if (dx > 0)
		{
			uint32_t next_id = (dx - 1) + board_size * dy;
			if(dfs_discovered.find(next_id) == dfs_discovered.end())
				dfs_stack.push_back(next_id);
		}
		if (dx < board_size - 1)
		{
			uint32_t next_id = (dx + 1) + board_size * dy;
			if (dfs_discovered.find(next_id) == dfs_discovered.end())
				dfs_stack.push_back(next_id);
		}
		if (dy > 0)
		{
			uint32_t next_id = dx + board_size * (dy - 1);
			if (dfs_discovered.find(next_id) == dfs_discovered.end())
				dfs_stack.push_back(next_id);
		}
		if (dy < board_size - 1)
		{
			uint32_t next_id = dx + board_size * (dy + 1);
			if (dfs_discovered.find(next_id) == dfs_discovered.end())
				dfs_stack.push_back(next_id);
		}
	}	
}

void EvaluateClusters()
{
	std::map<uint32_t, Cluster>::iterator cluster_map_it;
	cluster_map_it = cluster_map.begin();
	for (; cluster_map_it != cluster_map.end(); ++cluster_map_it)
	{
		std::set<uint32_t>::iterator cluster_tile_it, cluster_tile_it_end;
		cluster_tile_it = cluster_map_it->second.node_ids.begin();
		cluster_tile_it_end = cluster_map_it->second.node_ids.end();
		while (cluster_tile_it != cluster_tile_it_end)
		{
			Tile& t = board[*cluster_tile_it];
			std::set<uint32_t>::iterator tile_cluster_it;
			tile_cluster_it = t.clusters.begin();
			while(tile_cluster_it != t.clusters.end())
			{
				const Cluster& c = cluster_map[*tile_cluster_it];
				++tile_cluster_it;
				if (c.color == t.color)
					t.clusters.erase(std::prev(tile_cluster_it));				
			}
			++cluster_tile_it;
		}
		cluster_map_it->second.evaluate();
	}
}

bool SelectColor(uint32_t color)
{
	cluster_map.begin()->second.SetColor(color);
	std::set<uint32_t>::iterator it;
	it = cluster_map.begin()->second.node_ids.begin();
	for (; it != cluster_map.begin()->second.node_ids.end(); ++it)
	{
		uint32_t index = *it;
		uint32_t dx = index % board_size;
		uint32_t dy = index / board_size;
		ClusterTile(dx, dy);
	}
	EvaluateClusters();
	return cluster_map.begin()->second.node_ids.size() == board_size * board_size;
}

void ClusterEntireBoard()
{	
	for (uint32_t y = 0; y < board_size; ++y)
		for (uint32_t x = 0; x < board_size; ++x)
			ClusterTile(x, y);
}

void FillBoard()
{
	for(uint32_t y = 0; y < board_size; ++y)
		for (uint32_t x = 0; x < board_size; ++x)
		{
			auto seed = std::chrono::high_resolution_clock::now().time_since_epoch().count();
			std::mt19937 generator(seed);
			std::uniform_real_distribution<float> distribution(0.f, static_cast<float>(colors - 1));
			float rand = distribution(generator);
			uint32_t color = static_cast<uint32_t>(std::round(rand));				
			board[x + board_size * y].color = color;
		}
}

void GenerateColors()
{
	for (uint32_t c = 0; c < colors; ++c)
	{
		auto seed = std::chrono::high_resolution_clock::now().time_since_epoch().count();		
		std::mt19937 generator(seed);
		std::uniform_real_distribution<float> distribution(0x0, 0xFFFFFF);
		float rand = distribution(generator);
		ID2D1SolidColorBrush* brush;
		int32_t colorh = static_cast<int32_t>(rand);
		target->CreateSolidColorBrush(D2D1::ColorF(colorh, 1.0f), &brush);		
		//printf("Generated Color (%d): 0x%x\n", c, colorh);
		brushes.push_back(brush);
	}
}

typedef bool (*AlgorithmFunc)(void);

bool Algorithm_IterateColors()
{
	uint32_t pool_color = cluster_map.begin()->second.color;
	uint32_t next_color = pool_color >= colors - 1 ? 0 : pool_color + 1;	
	return SelectColor(next_color);
}

uint32_t alternate_index = 0;
uint32_t alternate_iter = 0;
bool alternate_direction = false;

void Reset_AlternateColors()
{
	alternate_index = 0;
	alternate_iter = 0;
	alternate_direction = false; //true = reverse
}
bool Algorithm_AlternateColors()
{	
	uint32_t pool_color = cluster_map.begin()->second.color;
	uint32_t next_color = pool_color;
	if (!alternate_direction)
	{
		next_color = alternate_index;
		alternate_direction = ++alternate_index == colors - 1;
		if (alternate_direction)
		{
			if (++alternate_iter == colors - 1)
				alternate_iter = 0;
		}
	}
	else
	{
		next_color = alternate_index;
		alternate_direction = !(--alternate_index == 0);
		if (!alternate_direction) alternate_index = alternate_iter;
	}
	return SelectColor(next_color);
}

bool Algorithm_HeaviestTreeCluster()
{
	uint32_t best_cluster_index = cluster_map.size();
	std::set<uint32_t>::iterator it;
	it = cluster_map.begin()->second.node_ids.begin();
	for (; it != cluster_map.begin()->second.node_ids.end(); ++it)
	{
		Tile& pool_tile = board[(*it)];
		std::set<uint32_t>::const_iterator pool_cluster_it;
		pool_cluster_it = pool_tile.clusters.begin();
		for (; pool_cluster_it != pool_tile.clusters.end(); ++pool_cluster_it)
		{
			uint32_t candidate_cluster_index = *pool_cluster_it;
			Cluster candidate_cluster = cluster_map[candidate_cluster_index];
			uint32_t candidate_cluster_size = candidate_cluster.clusters.size();
			if (best_cluster_index == cluster_map.size())
			{
				best_cluster_index = candidate_cluster_index;
				continue;
			}
			if (candidate_cluster_size > cluster_map[best_cluster_index].clusters.size())
			{
				best_cluster_index = candidate_cluster_index;
			}
		}
	}
	const Cluster& c = cluster_map[best_cluster_index];
	return SelectColor(c.color);
}

void ResetBoard()
{
	if (board != nullptr) delete[] board;
	board = new Tile[board_size * board_size];
	cluster_map.clear();
	cluster_id = 1;
	FillBoard();
	ClusterEntireBoard();
	EvaluateClusters();
}

RECT ClientRect;
int APIENTRY wWinMain(_In_ HINSTANCE hInstance, _In_opt_ HINSTANCE hPrevInstance, _In_ LPWSTR lpCmdLine, _In_ int nCmdShow)
{
	UNREFERENCED_PARAMETER(hPrevInstance);
	UNREFERENCED_PARAMETER(lpCmdLine);

	AllocConsole();
	FILE* p_file;
	freopen_s(&p_file, "CONIN$", "r", stdin);
	freopen_s(&p_file, "CONOUT$", "w", stdout);
	freopen_s(&p_file, "CONOUT$", "w", stderr);
	printf("Flood-It Solver by Aleksander Krimsky\n");
	
	std::string input;
	board_size_lbl:
	printf("Please enter a board size between 3 and 64: ");
	std::getline(std::cin, input);
	if (!SanitizeInput(input, 3, 64, board_size))
		goto board_size_lbl;

	colors_lbl:
	printf("Please enter the amount of colors (2 - 32): ");
	std::getline(std::cin, input);
	if (!SanitizeInput(input, 2, 32, colors))
		goto colors_lbl;

	cell_lbl:
	printf("Please enter the cell size px (default=16): ");
	std::getline(std::cin, input);
	if (!SanitizeInput(input, 8, 128, cell_size, 16))
		goto cell_lbl;

	sleep_lbl:
	printf("Please enter the sleep time ms (default=500): ");
	std::getline(std::cin, input);
	if (!SanitizeInput(input, 1, 10000, sleep_cycle, 500))
		goto sleep_lbl;

	printf("Press F1 to display the menu and enter a command.\n");

	uint32_t width = board_size * cell_size;
	uint32_t height = board_size * cell_size;
	ResetBoard();

	ClientRect.top = 0;
	ClientRect.left = 0;
	ClientRect.bottom = width;
	ClientRect.right = height;

	DWORD grfStyle, grfExStyle;
	grfStyle = WS_VISIBLE | WS_CLIPCHILDREN | WS_BORDER | WS_CAPTION | WS_SYSMENU | WS_MINIMIZEBOX;
	grfExStyle = WS_EX_STATICEDGE;
	AdjustWindowRectEx(&ClientRect, grfStyle, FALSE, grfExStyle);

	WNDCLASSEX MainClass = {};
	MainClass.cbClsExtra = 0;
	MainClass.cbWndExtra = 0;
	MainClass.lpszMenuName = 0;
	MainClass.hCursor = LoadCursor(NULL, IDC_ARROW);
	MainClass.hbrBackground = reinterpret_cast<HBRUSH>(COLOR_WINDOW + 1);
	MainClass.cbSize = sizeof(WNDCLASSEX);
	MainClass.hInstance = hInstance;
	MainClass.lpfnWndProc = WndProc;
	MainClass.lpszClassName = "Flood-It Solver Visualizer";
	MainClass.style = CS_HREDRAW | CS_VREDRAW;
	RegisterClassEx(&MainClass);

	AdjustWindowRect(&ClientRect, grfExStyle, FALSE);
	hWnd = CreateWindowEx(
		grfExStyle,
		MainClass.lpszClassName, "Flood-It Solver Visualizer",
		grfStyle,
		0, 0, 
		ClientRect.right - ClientRect.left, 
		ClientRect.bottom - ClientRect.top, 
		NULL, NULL, hInstance, NULL);

	RECT rc;
	GetClientRect(hWnd, &rc);
	D2D1_SIZE_U size = D2D1::SizeU(
		rc.right - rc.left,
		rc.bottom - rc.top
	);

	D2D1CreateFactory(D2D1_FACTORY_TYPE_MULTI_THREADED, &factory);
	factory->CreateHwndRenderTarget(
		D2D1::RenderTargetProperties(D2D1_RENDER_TARGET_TYPE_HARDWARE),
		D2D1::HwndRenderTargetProperties(hWnd, size, D2D1_PRESENT_OPTIONS_IMMEDIATELY),
		&target);

	target->SetAntialiasMode(D2D1_ANTIALIAS_MODE_ALIASED);
	GenerateColors();

	MSG msg;
	bool board_complete = false;
	bool run_simulation = false;
	uint32_t simulation_turns = 0;
	uint32_t max_turns = 0;
	uint32_t simulations_remaining = 0;
	uint32_t simulations = 0;
	float result_sum = 0.f;
	AlgorithmFunc algorithm = Algorithm_HeaviestTreeCluster;
	while (GetMessage(&msg, nullptr, 0, 0))
	{
		if (run_simulation)
		{
			simulation_turns++;
			board_complete = algorithm();
			run_simulation = !board_complete;
			if (!run_simulation)
			{
				result_sum += static_cast<float>(simulation_turns);
				printf("Finished running simulation in %d turns\n", simulation_turns);
				if (simulation_turns > max_turns) max_turns = simulation_turns;
				simulation_turns = 0;
				simulations_remaining--;
				if (simulations_remaining)
				{
					ResetBoard();
					Reset_AlternateColors();
					run_simulation = true;
				}
				else
				{
					float average = result_sum / static_cast<float>(simulations);
					printf("Ran %d simulations with an average of %f turns\n", simulations, average);
					printf("Board Size=%dx%d, Colors=%d\n", board_size, board_size, colors);
					printf("Worst Turns to solve a board: %d\n", max_turns);
				}
			}
			std::this_thread::sleep_for(std::chrono::milliseconds(sleep_cycle));
		}
		else if (GetAsyncKeyState(VK_F1) & 0x01)
		{
			printf("1. Generate New Board\n");
			printf("2. Run Algorithm Once\n");
			printf("3. Query Tile\n");
			printf("4. Query Cluster\n");
			printf("5. Run Complete Simulation\n");
			printf("Enter your selection:");
			uint32_t selection;
			selection_lbl:
			std::getline(std::cin, input);
			if (!SanitizeInput(input, 1, 5, selection))
				goto selection_lbl;
			switch (selection)
			{
			case 1:
				ResetBoard();
				run_simulation = false;
				break;
			case 2:
				if (!board_complete)
					board_complete = algorithm();
				else
					printf("The board is already completed\n");
				break;
			case 3:
			{
				printf("Enter the tile to query in zero-indexed \"x,y\" format:");
				std::string query_string_xy;
				std::getline(std::cin, query_string_xy);
				auto query_string_xy_it = query_string_xy.find(',');
				if (query_string_xy_it == std::string::npos)
				{
					printf("Missing comma delimiter\n");
					break;
				}
				std::string sx, sy;
				uint32_t x, y;
				sx = query_string_xy.substr(0, query_string_xy_it);
				query_string_xy.erase(0, query_string_xy_it + 1);
				sy = query_string_xy;
				if (!SanitizeInput(sx, 0, board_size - 1, x)) break;
				if (!SanitizeInput(sy, 0, board_size - 1, y)) break;
				const Tile& lt = board[x + board_size * y];
				printf("(%d, %d) Cluster=%d, Color=%d, Neighbors=%d\n", x, y,
					lt.cluster, lt.color, lt.clusters.size());				
			}
			break;
			case 4:
			{
				printf("Enter a cluster id:");
				std::string cluster_string;
				int cluster_id;
				std::getline(std::cin, cluster_string);
				cluster_id = atoi(cluster_string.c_str());
				std::map<uint32_t, Cluster>::const_iterator it;
				it = cluster_map.find(cluster_id);
				if (it == cluster_map.end())
				{
					printf("Could not find a cluster with id: %d\n", cluster_id);
					break;
				}
				printf("Color=%d, Size=%d, Neighbors=%d\n", it->second.color, 
					it->second.node_ids.size(), it->second.clusters.size());
			}
			break;
			case 5:
			{
				Reset_AlternateColors();
				simulations_lbl:
				printf("Please enter the amount of simulations to run:");
				std::getline(std::cin, input);
				if (!SanitizeInput(input, 1, 0xFFFFFF, simulations_remaining))
					goto simulations_lbl;
				run_simulation = simulations_remaining > 0;
				simulations = simulations_remaining;
				max_turns = 0;
			}
			break;
			}
			printf("Press F1 to display the menu and enter a command.\n");
		}
		TranslateMessage(&msg);
		DispatchMessage(&msg);
	}
	return (int)msg.wParam;
}

LRESULT CALLBACK WndProc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam)
{
	switch (message)
	{
	case WM_ERASEBKGND:
		break;
	case WM_PAINT:
	{
		target->BeginDraw();
		target->SetTransform(D2D1::Matrix3x2F::Identity());
		target->Clear();
		for (int32_t y = 0; y < board_size; ++y)
		{
			for (int32_t x = 0; x < board_size; ++x)
			{
				const Tile& tile = board[x + board_size * y];
				uint32_t color = tile.color;	
				D2D1_RECT_F tile_rect = D2D1::RectF(
					x * cell_size,
					y * cell_size,
					(x + 1) * cell_size,
					(y + 1) * cell_size);
				target->FillRectangle(&tile_rect, brushes.at(color));
			}
		}	
		target->EndDraw();
	}
	break;
	case WM_DESTROY:
		PostQuitMessage(0);
		break;
	default:
		return DefWindowProc(hWnd, message, wParam, lParam);
	}
	return 0;
}