from collections import deque
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_dijkstra_path(height_map: Grid, start: Coordinate, end: Coordinate) -> int | None:
    f_costs = []
    openNodes = {}
    closedNodes = {}

    openNodes[start] = (0, start.getDistanceTo(end, algorithm=DistanceAlgorithm.MANHATTAN), None)
    heappush(f_costs, (0, start))

    while f_costs:
        _, currentCoord = heappop(f_costs)
        if currentCoord not in openNodes:
            continue
        currentNode = openNodes[currentCoord]

        closedNodes[currentCoord] = currentNode
        del openNodes[currentCoord]
        if currentCoord == end:
            break

        for neighbour in height_map.getNeighboursOf(currentCoord, includeDiagonal=False):
            if neighbour in closedNodes or height_map.get(neighbour) - 1 > height_map.get(currentCoord):
                continue

            targetDist = neighbour.getDistanceTo(end, algorithm=DistanceAlgorithm.MANHATTAN)
            f_cost = targetDist + 1 + currentNode[1]

            if neighbour not in openNodes or f_cost < openNodes[neighbour][0]:
                openNodes[neighbour] = (f_cost, currentNode[1] + 1, currentCoord)
                heappush(f_costs, (f_cost, neighbour))

    if end not in closedNodes:
        return None
    else:
        steps = 0
        currentNode = closedNodes[end]
        while currentNode[2]:
            steps += 1
            currentNode = closedNodes[currentNode[2]]

        return steps


def get_bfs_path(height_map: Grid, end: Coordinate) -> int:
    v = {end: None}
    q = deque()
    q.append(end)
    while q:
        cur = q.popleft()
        cur_v = height_map.get(cur)
        if cur_v == 0:
            end = cur
            break

        for c in height_map.getNeighboursOf(cur, includeDiagonal=False):
            if c in v or height_map.get(c) + 1 < cur_v:
                continue

            v[c] = cur
            q.append(c)

    steps = 0
    while end in v:
        steps += 1
        end = v[end]

    return steps


class Day(AOCDay):
    inputs = [
        [
            (31, "input12_test"),
            (468, "input12_dennis"),
            (412, "input12_madEnrico"),
            (423, "input12"),
        ],
        [
            (29, "input12_test"),
            (459, "input12_dennis"),
            (402, "input12_madEnrico"),
            (416, "input12"),
        ]
    ]

    def get_map(self) -> (Grid, Coordinate, Coordinate):
        grid = Grid(30)
        start = None
        end = None
        for y, line in enumerate(self.getInput()):
            for x, char in enumerate(line):
                if char == 'S':
                    start = Coordinate(x, y)
                    grid.set(start, 0)
                elif char == 'E':
                    end = Coordinate(x, y)
                    grid.set(end, 25)
                else:
                    grid.set(Coordinate(x, y), ord(char) - ord('a'))

        return grid, start, end

    def part1(self) -> Any:
        return get_dijkstra_path(*self.get_map())

    def part2(self) -> Any:
        height_map, _, end = self.get_map()
        return get_bfs_path(height_map, end) - 1


if __name__ == '__main__':
    day = Day(2022, 12)
    day.run(verbose=True)