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(grid: Grid, start: Coordinate) -> int:
    q = deque([start])
    seen = set()
    nines = set()
    while q:
        pos = q.popleft()
        cur_val = grid.get(pos)
        if cur_val == 9:
            nines.add(pos)
            continue
        if pos in seen:
            continue
        seen.add(pos)

        for next_pos in grid.getNeighboursOf(pos, includeDiagonal=False):
            if grid.get(next_pos) == cur_val + 1:
                q.append(next_pos)

    return len(nines)


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


class Day(AOCDay):
    inputs = [
        [
            (1, "input10_test2"),
            (2, "input10_test3"),
            (4, "input10_test4"),
            (3, "input10_test5"),
            (36, "input10_test"),
            (644, "input10"),
        ],
        [
            (3, "input10_test7"),
            (81, "input10_test"),
            (1366, "input10"),
        ],
    ]

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

    def part1(self) -> Any:
        grid = self.parse_input()
        trail_heads = [
            Coordinate(x, y) for y in grid.rangeY() for x in grid.rangeX() if grid.get(Coordinate(x, y)) == 0
        ]
        return sum(find_nine(grid, trail) for trail in trail_heads)

    def part2(self) -> Any:
        grid = self.parse_input()
        trail_heads = [
            Coordinate(x, y) for y in grid.rangeY() for x in grid.rangeX() if grid.get(Coordinate(x, y)) == 0
        ]
        return sum(find_nine_ways(grid, trail) for trail in trail_heads)


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