#include "allocator.h"
#include "hmap.h"
#include "hmap_s.h"
#include "aoc.h"
#include "vec_s.h"
#include "dvec_s.h"
#include "util.h"

#include <assert.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>

struct layer {
	char cells[5][5];
	uint8_t adj[5][5];
};

struct cell {
	size_t adj;
	char c;
};

static void
load_map(struct hmap *map)
{
	const char *lpos, *lend;
	struct layer *layer;
	size_t x, y, len;
	char line[256];

	y = 0;
	layer = malloc(sizeof(struct layer));
	lpos = aoc.input;
	lend = lpos + aoc.input_size;
	while (readtok(line, sizeof(line), '\n', &lpos, lend)) {
		len = strlen(line);
		for (x = 0; x < len; x++)
			layer->cells[y][x] = line[x];
		y += 1;
	}

	hmap_add(map, (struct hmap_key) {.ss = 0},
		(struct hmap_val) {.p = layer});
}

static void
print_layer(struct hmap *map, ssize_t z)
{
	struct hmap_link *link;
	const struct layer *layer;
	size_t x, y;

	link = hmap_get(map, (struct hmap_key) {.ss = z});
	if (link) layer = link->value.p;
	for (y = 0; y < 5; y++) {
		for (x = 0; x < 5; x++) {
			if (x == 2 && y == 2) {
				aoc_debug("?");
			} else if (link) {
				aoc_debug("%c", layer->cells[y][x]);
			} else {
				aoc_debug(".");
			}
		}
		aoc_debug("\n");
	}
}

static void
print_layers(struct hmap *map, ssize_t minz, ssize_t maxz)
{
	ssize_t z;

	for (z = minz; z <= maxz; z++) {
		aoc_debug("Z %li\n", z);
		print_layer(map, z);
		aoc_debug("\n");
	}
}

static size_t
map_count(struct hmap *map)
{
	const struct layer *layer;
	struct hmap_iter iter;
	size_t x, y, count;

	count = 0;
	for (HMAP_ITER(map, iter)) {
		layer = iter.link->value.p;
		for (y = 0; y < 5; y++) {
			for (x = 0; x < 5; x++) {
				if (x == 2 && y == 2)
					continue;
				if (layer->cells[y][x] == '#')
					count += 1;
			}
		}
	}

	return count;
}

static size_t
top_cells(const struct layer *layer)
{
	size_t i, count;

	count = 0;
	for (i = 0; i < 5; i++)
		count += layer->cells[0][i] == '#';

	return count;
}

static size_t
right_cells(const struct layer *layer)
{
	size_t i, count;

	count = 0;
	for (i = 0; i < 5; i++)
		count += layer->cells[i][4] == '#';

	return count;
}

static size_t
bottom_cells(const struct layer *layer)
{
	size_t i, count;

	count = 0;
	for (i = 0; i < 5; i++)
		count += layer->cells[4][i] == '#';

	return count;
}

static size_t
left_cells(const struct layer *layer)
{
	size_t i, count;

	count = 0;
	for (i = 0; i < 5; i++)
		count += layer->cells[i][0] == '#';

	return count;
}

static size_t
cell_count(struct hmap *map, struct layer *layer,
	struct vec2i *pos, size_t dir, ssize_t z)
{
	struct hmap_link *link;
	const struct layer *nlayer;
	size_t count;

	count = 0;
	if (pos->x == 2 && pos->y == 2) {
		link = hmap_get(map, (struct hmap_key) {.ss = z + 1});
		if (!link) return 0;
		nlayer = link->value.p;
		switch (dir) {
		case ADJ_NORTH:
			return bottom_cells(nlayer);
		case ADJ_EAST:
			return left_cells(nlayer);
		case ADJ_SOUTH:
			return top_cells(nlayer);
		case ADJ_WEST:
			return right_cells(nlayer);
		}
		return count;
	} else if (pos->x < 0 || pos->x >= 5 || pos->y < 0 || pos->y >= 5) {
		link = hmap_get(map, (struct hmap_key) {.ss = z - 1});
		if (!link) return 0;
		nlayer = link->value.p;
		if (pos->x < 0) {
			return nlayer->cells[2][1] == '#';
		} else if (pos->x >= 5) {
			return nlayer->cells[2][3] == '#';
		} else if (pos->y < 0) {
			return nlayer->cells[1][2] == '#';
		} else if (pos->y >= 5) {
			return nlayer->cells[3][2] == '#';
		}
	} else {
		return layer->cells[pos->y][pos->x] == '#';
	}

	assert(0);
}

static void
step(struct hmap *map)
{
	struct hmap_iter iter;
	struct hmap_link **linkp;
	struct hmap_link *link;
	struct vec2i pos, npos;
	ssize_t z, minz, maxz;
	struct layer *layer;
	size_t i, count;
	char *cell;

	minz = maxz = 0;
	for (HMAP_ITER(map, iter)) {
		layer = iter.link->value._p;
		if (memchr(layer->cells, '#', 5 * 5)) {
			minz = MIN(minz, iter.link->key.ss);
			maxz = MAX(maxz, iter.link->key.ss);
		}
	}
	minz -= 1;
	maxz += 1;

	for (z = minz; z <= maxz; z++) {
		linkp = hmap_link_pos(map,
			(struct hmap_key) {.ss = z});
		if (!*linkp) {
			layer = malloc(sizeof(struct layer));
			memset(layer->cells, '.', 5 * 5);
			*linkp = hmap_link_alloc(map,
				(struct hmap_key) {.ss = z},
				(struct hmap_val) {.p = layer}, NULL);
		} else {
			layer = (*linkp)->value._p;
		}

		for (pos.y = 0; pos.y < 5; pos.y++) {
			for (pos.x = 0; pos.x < 5; pos.x++) {
				count = 0;
				for (i = 0; i < 4; i++) {
					vec2i_add(&npos, &pos, &adj[i]);
					count += cell_count(map,
						layer, &npos, i, z);
				}
				layer->adj[pos.y][pos.x] = (uint8_t) count;
			}
		}
	}

	for (z = minz; z <= maxz; z++) {
		for (pos.y = 0; pos.y < 5; pos.y++) {
			for (pos.x = 0; pos.x < 5; pos.x++) {
				link = hmap_get(map,
					(struct hmap_key) {.ss = z});
				layer = link->value._p;
				cell = &layer->cells[pos.y][pos.x];
				count = layer->adj[pos.y][pos.x];
				if (*cell == '#' && count != 1) {
					*cell = '.';
				} else if (*cell == '.' && (count == 1 || count == 2)) {
					*cell = '#';
				}
			}
		}
	}
}


void
part2(void)
{
	struct hmap_iter iter;
	struct hmap map;
	size_t i, answer;

	hmap_init(&map, 32, ssizet_hmap_hash, ssizet_hmap_keycmp,
		&stdlib_strict_heap_allocator);
	load_map(&map);

	for (i = 0; i < 200; i++) {
		step(&map);
		aoc_debug("-- MIN %lu --\n\n", i);
		print_layers(&map, -5, 5);
	}

	answer = map_count(&map);
	aoc.answer = aprintf("%lu", answer);
	aoc.solution = "1985";

	for (HMAP_ITER(&map, iter))
		free(iter.link->value._p);
	hmap_deinit(&map);
}