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

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)):
            if ny == y:
                continue
            yield (x, ny)

        for nx in list(range(max(0, x - 1), min(len(field[0]) - 1, x + 1) + 1)):
            if nx == x:
                continue
            yield (nx, y)

    @classmethod
    def lowest_points(cls, field: Field) -> Iterator[Coordinate]:
        for y, row in enumerate(field):
            for x, item in enumerate(row):
                if all(field[ny][nx] > item for nx, ny in cls.neighbours(field, x, y)):
                    yield x, y


class AssignmentOne(Assignment):
    example_result = 15

    def run(self, input: Field) -> int:
        lows = list(input[y][x] for x, y in self.lowest_points(input))
        return sum(lows) + len(lows)


def flatten(l: List[Set[Coordinate]]) -> List:
    flat_list = []
    for _ in l:
        flat_list += _

    return flat_list


class AssignmentTwo(Assignment):
    example_result = 1134

    @classmethod
    def find_basin_members(cls, field: Field, x: int, y: int) -> Set[Coordinate]:
        return {
            (x, y),
            *flatten([
                cls.find_basin_members(field, nx, ny)
                for nx, ny in cls.neighbours(field, x, y)
                if field[y][x] < field[ny][nx] < 9
            ])
        }

    def run(self, input: Field) -> int:
        return reduce(
            mul,
            sorted([
                len(self.find_basin_members(input, x, y))
                for x, y in self.lowest_points(input)
            ], reverse=True)[:3],
            1
        )