from collections import defaultdict

import aoclib

DAY = 17
TEST_SOLUTION_PART1 = 112
TEST_SOLUTION_PART2 = 848


def simulate(active_cubes, fourD=False):
    active_neighbours = defaultdict(int)
    for active_cube in active_cubes:
        if fourD:
            x, y, z, w = active_cube
        else:
            x, y, z = active_cube

        for cx in [x - 1, x, x + 1]:
            for cy in [y - 1, y, y + 1]:
                for cz in [z - 1, z, z + 1]:
                    if not fourD:
                        if not (x == cx and y == cy and z == cz):
                            active_neighbours[(cx, cy, cz)] += 1
                    else:
                        for cw in [w - 1, w, w + 1]:
                            if not (x == cx and y == cy and z == cz and w == cw):
                                active_neighbours[(cx, cy, cz, cw)] += 1

    new_active_cubes = set()
    for coord, neighbour_count in active_neighbours.items():
        if coord in active_cubes and neighbour_count in (2, 3):
            new_active_cubes.add(coord)
        elif neighbour_count == 3:
            new_active_cubes.add(coord)

    return new_active_cubes


def get_active_cubes_from_input(grid_input, dimensions=3):
    active_cubes = set()
    for x, line in enumerate(grid_input):
        for y, char in enumerate(line):
            if char == '#':
                active_cubes.add((x, y) + tuple(0 for _ in range(dimensions - 2)))

    return active_cubes


def part1(test_mode=False):
    my_input = aoclib.getInputAsArray(day=DAY, test=test_mode)
    active_cubes = get_active_cubes_from_input(my_input)

    for _ in range(6):
        active_cubes = simulate(active_cubes)

    return len(active_cubes)


def part2(test_mode=False):
    my_input = aoclib.getInputAsArray(day=DAY, test=test_mode)
    active_cubes = get_active_cubes_from_input(my_input, 4)
    for _ in range(6):
        active_cubes = simulate(active_cubes, True)

    return len(active_cubes)