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 walk(grid: Grid, start: Coordinate, min_steps: int = 0, max_steps: int = 3) -> int:
    q = [(0, start, Coordinate(0, 0), 0)]
    v = {}
    while q:
        cost, pos, last_dir, last_dir_count = heappop(q)
        if (pos, last_dir, last_dir_count) in v:
            continue
        v[pos, last_dir, last_dir_count] = cost
        if 0 < last_dir_count < min_steps:
            n = pos + last_dir
            if grid.isWithinBoundaries(n):
                heappush(q, (cost + grid.get(n), n, last_dir, last_dir_count + 1))
        else:
            for n in grid.getNeighboursOf(pos, includeDiagonal=False, includeDefault=False):
                new_dir = n - pos
                new_dir_count = last_dir_count + 1 if new_dir == last_dir else 1
                if new_dir.reverse() != last_dir and (new_dir != last_dir or last_dir_count < max_steps):
                    heappush(q, (cost + grid.get(n), n, new_dir, new_dir_count))

    min_cost = int(1e9)
    for (pos, _, _), cost in v.items():
        if pos == Coordinate(grid.maxX, grid.maxY):
            min_cost = min(min_cost, cost)

    return min_cost


class Day(AOCDay):
    inputs = [
        [
            (102, "input17_test"),
            (870, "input17"),
        ],
        [
            (94, "input17_test"),
            (1063, "input17"),
        ]
    ]

    def part1(self) -> Any:
        grid = Grid.from_data(self.getInput(), translate={'[0-9]': int})
        return walk(grid, Coordinate(0, 0))

    def part2(self) -> Any:
        grid = Grid.from_data(self.getInput(), translate={'[0-9]': int})
        return walk(grid, Coordinate(0, 0), 4, 10)


if __name__ == '__main__':
    day = Day(2023, 17)
    day.run(verbose=True)