from collections import deque

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


def find_nine_ways(grid: Grid, start: Coordinate, next_val: int = 1) -> list[Coordinate]:
    if next_val == 9:
        return [x for x in grid.getNeighboursOf(start, includeDiagonal=False) if grid.get(x) == 9]
    else:
        ret = []
        for x in grid.getNeighboursOf(start, includeDiagonal=False):
            if grid.get(x) == next_val:
                ret += find_nine_ways(grid, x, next_val + 1)

        return ret

class Day(AOCDay):
    inputs = [
        [
            (36, "input10_test"),
            (644, "input10"),
        ],
        [
            (81, "input10_test"),
            (1366, "input10"),
        ],
    ]

    def parse_input(self) -> Grid:
        return Grid.from_data(self.getInput(), default=-1, translate={"[0-9]": int})

    def part1(self) -> Any:
        grid = self.parse_input()
        return sum(len(set(find_nine_ways(grid, trail))) for trail in grid.find(0))

    def part2(self) -> Any:
        grid = self.parse_input()
        return sum(len(find_nine_ways(grid, trail)) for trail in grid.find(0))


if __name__ == "__main__":
    day = Day(2024, 10)
    day.run(verbose=True)