from collections import defaultdict

import aoclib

DAY = 17
TEST_SOLUTION_PART1 = 112
TEST_SOLUTION_PART2 = 848


def getEdgeCoords(grid, quad=False):
    x_list = [x[0] for x in grid.keys()]
    y_list = [x[1] for x in grid.keys()]
    z_list = [x[2] for x in grid.keys()]
    if not quad:
        return min(x_list), min(y_list), min(z_list), max(x_list), max(y_list), max(z_list)
    else:
        w_list = [x[3] for x in grid.keys()]
        return min(x_list), min(y_list), min(z_list), min(w_list), max(x_list), max(y_list), max(z_list), max(w_list)


def simulate(grid):
    minX, minY, minZ, maxX, maxY, maxZ = getEdgeCoords(grid)

    new_grid = defaultdict(int)

    for x in range(minX - 1, maxX + 2):
        for y in range(minY - 1, maxY + 2):
            for z in range(minZ - 1, maxZ + 2):
                cube_state = grid[(x, y, z)]

                sum_neighbours = 0
                for cx in [x - 1, x, x + 1]:
                    for cy in [y - 1, y, y + 1]:
                        for cz in [z - 1, z, z + 1]:
                            sum_neighbours += grid[(cx, cy, cz)]

                sum_neighbours -= cube_state

                if not cube_state and sum_neighbours == 3:
                    new_grid[(x, y, z)] = 1
                elif cube_state:
                    if sum_neighbours not in [2, 3]:
                        new_grid[(x, y, z)] = 0
                    else:
                        new_grid[(x, y, z)] = 1

    return new_grid


def simulate4D(grid):
    minX, minY, minZ, minW, maxX, maxY, maxZ, maxW = getEdgeCoords(grid, True)

    new_grid = defaultdict(int)

    for x in range(minX - 1, maxX + 2):
        for y in range(minY - 1, maxY + 2):
            for z in range(minZ - 1, maxZ + 2):
                for w in range(minW - 1, maxW + 2):
                    cube_state = grid[(x, y, z, w)]

                    sum_neighbours = 0
                    for cx in [x - 1, x, x + 1]:
                        for cy in [y - 1, y, y + 1]:
                            for cz in [z - 1, z, z + 1]:
                                for cw in [w - 1, w, w + 1]:
                                    sum_neighbours += grid[(cx, cy, cz, cw)]

                    sum_neighbours -= cube_state

                    if not cube_state and sum_neighbours == 3:
                        new_grid[(x, y, z, w)] = 1
                    elif cube_state:
                        if sum_neighbours not in [2, 3]:
                            new_grid[(x, y, z, w)] = 0
                        else:
                            new_grid[(x, y, z, w)] = 1

    return new_grid


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

    grid = defaultdict(int)

    for y, line in enumerate(my_input):
        for x, char in enumerate(line):
            grid[(x, y, 0)] = 1 if char == '#' else 0

    for i in range(6):
        grid = simulate(grid)

    return sum(grid.values())


def part2(test_mode=False):
    my_input = aoclib.getInputAsArray(day=DAY, test=test_mode)

    grid = defaultdict(int)

    for y, line in enumerate(my_input):
        for x, char in enumerate(line):
            grid[(x, y, 0, 0)] = 1 if char == '#' else 0

    for i in range(6):
        grid = simulate4D(grid)

    return sum(grid.values())