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 x in range(caveMap.minX, caveMap.maxX + 1):
        for y in range(caveMap.minY, caveMap.maxY + 1):
            thisPoint = Coordinate(x, y)
            lowest = True

            for t in thisPoint.getNeighbours(includeDiagonal=False):
                if caveMap.get(t) <= caveMap.get(Coordinate(x, y)):
                    lowest = False

            if lowest:
                lowPoints.append(thisPoint)

    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 and caveMap.get(n) != 99:
            visited.add(n)
            visited = set.union(visited, getBasin(caveMap, n, visited))

    return visited


class Day(AOCDay):
    test_solutions_p1 = [15]
    test_solutions_p2 = [1134]

    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])