from collections import defaultdict
from utils import open_day, Point2D, adjacent_bounded
from heapq import heappop, heappush

def a_star(start, goal, neighbours, h, d):
    vtoi = dict()
    itov = dict()
    i = 0
    def intern(v):
        nonlocal i
        ret = vtoi.get(v)
        if ret is None:
            ret = i
            vtoi[v] = i
            itov[i] = v
            i += 1
        return ret
    open_set = {start}
    open_heapq = [(0, intern(start))]
    came_from = dict()
    g_score = defaultdict(lambda: float('inf'))
    g_score[start] = 0
    f_score = defaultdict(lambda: float('inf'))
    f_score[start] = h(start)
    while open_set:
        while True:
            f, current = heappop(open_heapq)
            current = itov[current]
            if current in open_set: break
        if current == goal:
            return f
        open_set.remove(current)
        for neighbour in neighbours(current):
            tentative_g = g_score[current] + d(current, neighbour)
            if tentative_g < g_score[neighbour]:
                came_from[neighbour] = current
                g_score[neighbour] = tentative_g
                f_score[neighbour] = tentative_g + h(neighbour)
                open_set.add(neighbour)
                heappush(open_heapq, (f_score[neighbour], intern(neighbour)))
    return None

with open_day(15) as f:
    inp = [[int(c) for c in l.rstrip()] for l in f]

def d(a, b):
    add = b.x // len(inp[0]) + b.y // len(inp)
    x, y = b.x % len(inp[0]), b.y % len(inp)
    return (inp[y][x] + add - 1) % 9 + 1

start = Point2D(0, 0)
goal1 = Point2D(len(inp[0]) - 1, len(inp) - 1)
neighbours1 = lambda p: adjacent_bounded(p, Point2D(len(inp[0]), len(inp)), False)
print(a_star(start, goal1, neighbours1, lambda p: 0, d))
goal2 = Point2D(len(inp[0]) * 5 - 1, len(inp) * 5 - 1)
neighbours2 = lambda p: adjacent_bounded(p, Point2D(len(inp[0]) * 5, len(inp) * 5), False)
print(a_star(start, goal2, neighbours2, lambda p: 0, d))