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


def getCaveMapFromInput(puzzle_input: List[str]) -> Grid:
    caveMap = Grid(99)
    for y, s in enumerate(puzzle_input):
        for x, v in enumerate(s):
            caveMap.set(Coordinate(x, y), int(v))

    return caveMap


def getLowPoints(caveMap: Grid) -> List[Coordinate]:
    lowPoints = []
    for point in caveMap.getActiveCells():
        if not sum((caveMap.get(t) <= caveMap.get(point)) for t in point.getNeighbours(includeDiagonal=False)):
            lowPoints.append(point)

    return lowPoints


def getBasin(caveMap: Grid, start: Coordinate, visited: set) -> set:
    for n in start.getNeighbours(includeDiagonal=False):
        if n in visited:
            continue

        if caveMap.get(n) < 9:
            visited.add(n)
            visited = set.union(visited, getBasin(caveMap, n, visited))

    return visited


class Day(AOCDay):
    inputs = [
        [
            (15, "test_input09"),
            (545, "input09")
        ],
        [
            (1134, "test_input09"),
            (950600, "input09")
        ]
    ]

    def part1(self) -> Any:
        caveMap = getCaveMapFromInput(self.getInput())
        return sum(caveMap.get(x) + 1 for x in getLowPoints(caveMap))

    def part2(self) -> Any:
        caveMap = getCaveMapFromInput(self.getInput())
        lowPoints = getLowPoints(caveMap)
        basins = list(sorted([getBasin(caveMap, point, set()) for point in lowPoints], key=lambda l: len(l)))

        return len(basins[-1]) * len(basins[-2]) * len(basins[-3])


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