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


def getOctopusGrid(lines: List[str]) -> Grid:
    grid = Grid(-1)
    for y, l in enumerate(lines):
        for x, c in enumerate(l):
            grid.set(Coordinate(x, y), int(c))

    return grid


def flashGrid(grid: Grid) -> int:
    flashers = set()
    for c in grid.getActiveCells():
        if grid.add(c) > 9:
            flashers.add(c)

    did_flash = set()
    while flashers:
        thisFlash = flashers.pop()
        did_flash.add(thisFlash)
        grid.set(thisFlash, 0)

        for n in grid.getNeighboursOf(thisFlash):
            if n not in did_flash and grid.add(n) > 9:
                flashers.add(n)

    return len(did_flash)


class Day(AOCDay):
    test_solutions_p1 = [1656]
    test_solutions_p2 = [195]

    def part1(self) -> Any:
        grid = getOctopusGrid(self.getInput())
        return sum([flashGrid(grid) for _ in range(100)])

    def part2(self) -> Any:
        grid = getOctopusGrid(self.getInput())
        step_count = 0
        while grid.getSum() != 0:
            step_count += 1
            flashGrid(grid)

        return step_count