import aoclib

DAY = 24
TEST_SOLUTION_PART1 = 10
TEST_SOLUTION_PART2 = 2208


def getTileGrid(from_input):
    grid = {}
    for line in from_input:
        x = 0
        y = 0
        while len(line) > 0:
            if line[0] == 'e':
                x += 1
                line = line[1:]
            elif line[0] == 'w':
                x -= 1
                line = line[1:]
            elif line[0] == 'n':
                y -= 1
                if line[1] == 'e':
                    x += 1
                line = line[2:]
            elif line[0] == 's':
                y += 1
                if line[1] == 'w':
                    x -= 1
                line = line[2:]

        if (x, y) in grid:
            del grid[(x, y)]
        else:
            grid[(x, y)] = True

    return grid


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

    grid = getTileGrid(my_input)

    return len(grid)


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

    for i in range(100):
        black_neighbour_count = {}
        for coord in grid:
            x, y = coord
            for nCoord in [(x-1, y), (x-1, y+1), (x, y-1), (x, y+1), (x+1, y-1), (x+1, y)]:
                black_neighbour_count[nCoord] = black_neighbour_count.get(nCoord, 0) + 1

        for nCoord, nCount in black_neighbour_count.items():
            if nCoord in grid and nCount > 2:
                del grid[nCoord]
            elif nCoord not in grid and nCount == 2:
                grid[nCoord] = True

        for coord in grid.copy():
            if coord not in black_neighbour_count:
                del grid[coord]

    return len(grid)