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


def part_1(input):
    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):
    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)