from abc import ABC
from collections import Counter
from copy import copy
from functools import reduce
from operator import mul
from typing import List, Tuple, Dict, FrozenSet, Iterator, Set, Callable, Any

from aoc import BaseAssignment

Coordinate = Tuple[int, int]
Field = List[List[int]]


class Assignment(BaseAssignment, ABC):
    def parse_item(self, item: str) -> List[int]:
        return [int(i) for i in item]

    def read_input(self, example = False) -> Field:
        return list(super().read_input(example))

    @classmethod
    def neighbours(cls, field: Field, x: int, y: int) -> Iterator[Coordinate]:
        for ny in list(range(max(0, y - 1), min(len(field) - 1, y + 1) + 1)):
            for nx in list(range(max(0, x - 1), min(len(field[0]) - 1, x + 1) + 1)):
                if ny == y and nx == x:
                    continue
                yield (nx, ny)

    @classmethod
    def flatten(cls, l: List[List[Any]]) -> List:
        flat_list = []
        for _ in l:
            flat_list += _

        return flat_list

    @classmethod
    def flash(cls, field, x, y):
        for nx, ny in cls.neighbours(field, x, y):
            if field[ny][nx] > 0:
                field[ny][nx] = (field[ny][nx] + 1) % 10

                if field[ny][nx] == 0:
                    cls.flash(field, nx, ny)

    @classmethod
    def step(cls, field: Field):
        flashing = []

        for y, row in enumerate(field):
            for x, item in enumerate(row):
                field[y][x] = (item + 1) % 10
                if field[y][x] == 0:
                    flashing.append((x, y))

        for x, y in flashing:
            cls.flash(field, x, y)


class AssignmentOne(Assignment):
    example_result = 1656

    def run(self, input: Field) -> int:
        count = 0
        for _ in range(100):
            self.step(input)
            count += Counter(self.flatten(input))[0]

        return count


class AssignmentTwo(Assignment):
    example_result = 195

    def run(self, input: Field) -> int:
        step_nr = 0
        while True:
            step_nr += 1

            self.step(input)

            flat_input = self.flatten(input)

            if len(flat_input) == Counter(flat_input)[0]:
                break

        return step_nr