from heapq import heappush, heappop
from tools.aoc import AOCDay
from tools.coordinate import Coordinate, DistanceAlgorithm
from tools.grid import Grid
from typing import Any


def get_path(memory: Grid) -> set[Coordinate]:
    queue = [(0, 0, Coordinate(0, 0), set())]
    target = Coordinate(memory.maxX, memory.maxY)
    visited = set()
    while queue:
        _, cost, pos, path = heappop(queue)
        if pos == target:
            return path
        if pos in visited:
            continue
        visited.add(pos)

        for n in memory.getNeighboursOf(pos, includeDiagonal=False, includeDefault=True):
            if memory.get(n):
                continue

            target_dist = n.getDistanceTo(target, algorithm=DistanceAlgorithm.MANHATTAN, includeDiagonals=False)
            heappush(queue, (cost + target_dist, cost + 1, n, path | {n}))

    return set()


class Day(AOCDay):
    inputs = [
        [
            (22, "input18_test"),
            (248, "input18"),
        ],
        [
            ("6,1", "input18_test"),
            ("32,55", "input18"),
        ],
    ]

    def parse_input(self) -> tuple[int, Grid, list[Coordinate]]:
        memory = Grid()
        memory.maxX = 70 if not self.is_test() else 6
        memory.maxY = 70 if not self.is_test() else 6
        falling_bytes = [Coordinate(x, y) for x, y in self.getIntsFromInput()]
        steps = 1024 if not self.is_test() else 12
        for i in range(steps):
            memory.set(falling_bytes[i])

        return steps, memory, falling_bytes

    def part1(self) -> Any:
        _, memory, falling_bytes = self.parse_input()
        return len(get_path(memory))

    def part2(self) -> Any:
        steps, memory, falling_bytes = self.parse_input()

        path = get_path(memory)
        for i in range(steps + 1, len(falling_bytes)):
            memory.set(falling_bytes[i])
            if falling_bytes[i] not in path:
                continue
            path = get_path(memory)
            if not path:
                return str(falling_bytes[i])[1:-1]


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