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


def get_best_score(grid: Grid) -> tuple[int, set[Coordinate]]:
    target = Coordinate(grid.maxX - 1, 1)
    queue = [(0, Coordinate(1, grid.maxY - 1), Coordinate(1, 0), {Coordinate(1, grid.maxY - 1)})]
    visited = set()
    while queue:
        score, pos, facing, path = heappop(queue)
        if pos == target:
            return score, path

        if pos in visited:
            continue
        visited.add(pos)

        for n in grid.getNeighboursOf(pos, includeDiagonal=False):
            if grid.get(n) or n in visited:
                continue

            add_turn = 1001 if n - pos != facing else 1
            heappush(queue, (score + add_turn, n, n - pos, path | {n}))


def get_tiles_from_all_paths(grid: Grid) -> int:
    target_score, target_path = get_best_score(grid)
    target = Coordinate(grid.maxX - 1, 1)
    queue = [(0, Coordinate(1, grid.maxY - 1), Coordinate(1, 0), {Coordinate(1, grid.maxY - 1)})]
    # all_paths = set()
    visited = {}
    while queue:
        score, pos, facing, path = heappop(queue)
        if pos in visited and score > visited[pos][0] + 1000:
            continue
        if pos in target_path and pos in visited:
            # print(f"{pos} in target_path and visited. {score=}, {visited[pos]=}, discarding {path=}")
            if score <= visited[pos][0] or score == visited[pos][0] + 1000:
                target_path |= path
            continue

        visited[pos] = (score, facing)
        if pos == target:
            target_path |= path
            continue

        if score >= target_score:
            continue

        for n in grid.getNeighboursOf(pos, includeDiagonal=False):
            if grid.get(n) or n in path:
                continue

            add_turn = 1001 if n - pos != facing else 1
            heappush(queue, (score + add_turn, n, n - pos, path | {n}))

    for p in target_path:
        grid.set(p, "O")
    grid.print()
    return len(target_path)


class Day(AOCDay):
    inputs = [
        [
            (7036, "input16_test"),
            (11048, "input16_test2"),
            (93436, "input16"),
        ],
        [
            (45, "input16_test"),
            (64, "input16_test2"),
            (486, "input16"),
        ],
    ]

    def parse_input(self) -> Grid:
        return Grid().from_data(self.getInput(), default=None, translate={r"[\.SE]": False, "#": True})

    def part1(self) -> Any:
        return get_best_score(self.parse_input())[0]

    def part2(self) -> Any:
        return get_tiles_from_all_paths(self.parse_input())


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