import itertools

# TODO: Should FixedGrid allow offset domains (say -1->4 for x and 3->8 for y)?
class FixedGrid:
    '''Fixed size grid utility class to simplify grid operations and minimize
    bugs.
    '''
    def __init__(self, grid):
        self._grid = grid
        self._width = len(self._grid)
        self._height = len(self._grid[0])

    @staticmethod
    def parse(s, linesplit_fn=None, line_separator='\n', value_fn=None):
        grid = []
        for line in s.split(line_separator):
            if linesplit_fn is not None:
                line = linesplit_fn(line)
            if value_fn is None:
                grid.append(list(line))
            else:
                grid.append(list(map(value_fn, line)))
        return FixedGrid(grid).transpose()

    @staticmethod
    def from_dict(d, missing=None):
        '''Converts a coordinate->value dictionary to a FixedGrid.
        Expects that minimum x and y coordinates in the grid are both 0.

        Arguments:
        missing -- Value to use for any coordinates not present in the dictionary
        '''
        low_x, high_x = None, None
        low_y, high_y = None, None
        for x, y in d.keys():
            if low_x is None:
                low_x, high_x = x, x
                low_y, high_y = y, y
            else:
                low_x = min(low_x, x)
                high_x = max(high_x, x)
                low_y = min(low_y, y)
                high_y = max(high_y, y)

        assert(low_x == low_y == 0)

        return FixedGrid([[d.get((x, y), missing)
                           for y in range(low_y, high_y+1)]
                          for x in range(low_x, high_x+1)])

    def to_dict(self):
        return {(x,y): val
                for x,col in enumerate(self._grid)
                for y,val in enumerate(col)}

    def transpose(self):
        return FixedGrid([[self._grid[x][y]
                           for x in range(self._width)]
                          for y in range(self._height)])

    @property
    def width(self):
        return self._width

    @property
    def height(self):
        return self._height

    @property
    def area(self):
        return self._width * self._height

    def __getitem__(self, c):
        x, y = c
        assert(0 <= x < self._width and 0 <= y < self._height)
        return self._grid[x][y]

    def __setitem__(self, c, val):
        x, y = c
        assert(0 <= x < self._width and 0 <= y < self._height)
        self._grid[x][y] = val

    def items(self, column_first = False):
        '''Generates all coordinate,value pairs in the grid for iteration.

        Arguments:
        column_first -- Iterate by column first rather than by row first
        '''
        if column_first:
            for y in range(self._height):
                for x, col in enumerate(self._grid):
                    yield (x, y), col[y]
        else:
            for x, col in enumerate(self._grid):
                for y, val in enumerate(col):
                    yield (x, y), val

    def neighbors(self, x, y, diagonals=False):
        assert(0 <= x < self._width and 0 <= y < self._height)
        if diagonals:
            for nx, ny in itertools.product((x-1, x, x+1),
                                            (y-1, y, y+1)):
                if x == nx and y == ny:
                    continue
                if 0 <= nx < self._width and 0 <= ny < self._height:
                    yield nx, ny
        else:
            if 0 < x:
                yield x-1, y
            if x+1 < self._width:
                yield x+1, y
            if 0 < y:
                yield x, y-1
            if y+1 < self._height:
                yield x, y+1

    def print(self, line_spacing=' '):
        print('\n'.join(line_spacing.join(str(self._grid[x][y])
                                          for x in range(self._width))
                        for y in range(self._height)))

# TODO: ExpandingGrid (probably dict-based)

def parse_input(s):
    algo, image = s.split('\n\n')
    image = FixedGrid.parse(image)
    return algo, image.to_dict(), image.width, image.height

def neighbors(x, y):
    for ay in (-1, 0, 1):
        for ax in (-1, 0, 1):
            yield (x+ax, y+ay)

def step(algo, image, xrange, yrange, default):
    new_im = {}
    xrange = range(xrange[0]-1, xrange[-1]+2)
    yrange = range(yrange[0]-1, yrange[-1]+2)
    for x in xrange:
        for y in yrange:
            section = [image.get(n, default)
                       for n in neighbors(x,y)]
            section = ''.join(section)
            section = section.replace('.', '0')
            section = section.replace('#', '1')
            idx = int(section, 2)
            new_im[x,y] = algo[idx]

    # Figure out new out of bounds default
    section = (default*9).replace('.', '0').replace('#', '1')
    default = algo[int(section, 2)]

    return new_im, xrange, yrange, default

def part1(s):
    algo, image, width, height = parse_input(s)
    xrange = range(width)
    yrange = range(height)

    default = '.'

    for _ in range(2):
        image, xrange, yrange, default = step(algo, image, xrange, yrange, default)

    answer = sum(1 for c in image.values() if c == '#')

    print(f'The answer to part one is {answer}')

def part2(s):
    algo, image, width, height = parse_input(s)
    xrange = range(width)
    yrange = range(height)

    default = '.'

    for _ in range(50):
        image, xrange, yrange, default = step(algo, image, xrange, yrange, default)

    answer = sum(1 for c in image.values() if c == '#')

    print(f'The answer to part two is {answer}')

with open('./2021/20.input') as input_file:

    INPUT = input_file.read()
part1(INPUT)
part2(INPUT)