#!/usr/bin/env python3
from pathlib import Path
from heapq import heappush, heappop


def part_1(input):
    """The solution for Part 1 is an implementation of dijkstra algorithm. The input data is stored
    in a dictionary 'nodes' where the keys are a tuple of the coordinates and the value is the cost
    (risk level) to travel to this node.
    The total cost for each node will be stored in a dictionary named 'cost'. We start with a queue,
    that will contain all the nodes we have to check. Initially there is only the starting node
    known. The queue also contains the cost for each node to be able to prioritize them in an heapq.
    The 'done' set is used to remember the already processed nodes.
    In each iteration we take the node with lowest cost known so far from the queue and set it as
    the active node. If this node is already stored in 'done' we can skip it and start the next
    iteration immediately. Nodes can be pushed to the queue multiple times but we only need to
    process them once, hence we always use the node with the lowest cost from the queue.
    If the node hasn't been processed yet we check all of its four neighbors if they have been
    processed. If not we calculate the total cost for this node by adding the risk level to the cost
    of the active node. The total cost will then be compared to the current known cost for this new
    node. If it is lower or not yet existing the cost will be updated in the 'cost' dictionary and
    the new node will be pushed to the queue. At this point we could also push the active node to a
    dictionary containing the previous node for each node to be able to construct the path. But
    since only the cost is asked we don't store this information.
    Finally the active node will be pushed to the 'done' set and start the next iteration if there
    are elements in the queue.
    """
    result = 0
    x_size = len(input[0].strip())
    y_size = len(input)
    nodes = {(x, y): int(v) for y, line in enumerate(input)
             for x, v in enumerate(line.strip())}
    start = (0, 0)
    end = (x_size - 1, y_size - 1)
    cost = {start: 0}
    queue = [(0, start)]
    done = set()
    neighbors = set([(-1, 0), (0, -1), (0, 1), (1, 0)])
    while len(queue):
        cur_cost, current = heappop(queue)
        if current in done:
            continue
        (xc, yc) = current
        next = set()
        for (xn, yn) in neighbors:
            nbr = (xc + xn, yc + yn)
            if nbr in nodes and not nbr in done:
                next.add(nbr)
        for n in next:
            next_cost = cur_cost + nodes[n]
            if not n in cost or next_cost < cost[n]:
                cost[n] = next_cost
                heappush(queue, (next_cost, n))
        done.add(current)
    result = cost[end]
    print("Part 1 result:", result)


def part_2(input):
    """The part 2 solution is essentially the same as the solution for part 1. The only difference
    is that the grid now repeats itself five times to the the right and downwards. The grid itself
    is only stored once and the new boundaries and costs for new nodes are calculated on the fly
    while the iteration across the nodes happen.
    """
    result = 0
    x_size = len(input[0].strip())
    y_size = len(input)
    repeat = 5
    nodes = {(x, y): int(v) for y, line in enumerate(input)
             for x, v in enumerate(line.strip())}
    start = (0, 0)
    end = ((x_size * repeat) - 1, (y_size * repeat) - 1)
    cost = {start: 0}
    queue = [(0, start)]
    done = set()
    neighbors = set([(-1, 0), (0, -1), (0, 1), (1, 0)])
    while len(queue):
        cur_cost, current = heappop(queue)
        if current in done:
            continue
        (xc, yc) = current
        next = set()
        for (xn, yn) in neighbors:
            x = xc + xn
            y = yc + yn
            if 0 <= x < (x_size * repeat) and 0 <= y < (y_size * repeat):
                if not (x, y) in done:
                    next.add((x, y))
        for (x, y) in next:
            next_node_cost = nodes[(x % x_size, y % y_size)] + \
                (x // x_size) + (y // y_size)
            if next_node_cost > 9:
                next_node_cost -= 9
            if not (x, y) in cost or (cur_cost + next_node_cost) < cost[(x, y)]:
                cost[(x, y)] = cur_cost + next_node_cost
                heappush(queue, (cur_cost + next_node_cost, (x, y)))
        done.add(current)
    result = cost[end]
    print("Part 2 result:", result)


input = list()
p = Path(__file__).with_name('input.txt')
with open(p) as f:
    input = f.readlines()
    part_1(input)
    part_2(input)