from tools.aoc import AOCDay
from tools.coordinate import Coordinate
from tools.grid import Grid
from typing import Any, List


def enhance(image: Grid, enhancements: List[bool], count: int) -> Grid:
    flippingInfinity = enhancements[0] and not enhancements[-1]
    infinity = False

    for _ in range(count):
        if flippingInfinity:
            infinity = not infinity

        new_image = Grid(infinity)
        for x in image.rangeX(pad=1):
            for y in image.rangeY(pad=1):
                index_str = ''
                for y1 in [y - 1, y, y + 1]:
                    for x1 in [x - 1, x, x + 1]:
                        index_str += '1' if image.get(Coordinate(x1, y1)) else '0'

                new_image.set(Coordinate(x, y), enhancements[int(index_str, 2)])

        image = new_image

    return image


class Day(AOCDay):
    inputs = [
        [
            (35, "test_input20"),
            (5425, "input20")
        ],
        [
            (3351, "test_input20"),
            (14052, "input20")
        ]
    ]

    def getInputs(self) -> (Grid, List[bool]):
        enhancement, image = self.getMultiLineInputAsArray()
        enhancement = [c == '#' for c in enhancement[0]]
        grid = Grid()
        for y, l in enumerate(image):
            for x, c in enumerate(l):
                grid.set(Coordinate(x, y), c == '#')

        return grid, enhancement

    def part1(self) -> Any:
        image, enhancement = self.getInputs()
        image = enhance(image, enhancement, 2)
        return image.getSum()

    def part2(self) -> Any:
        image, enhancement = self.getInputs()
        image = enhance(image, enhancement, 50)
        return image.getSum()


if __name__ == '__main__':
    day = Day(2021, 20)
    day.run(verbose=True)