#!/usr/bin/env python

import numpy as np
import math
from heapq import *
from collections import defaultdict

with open('day15.txt') as data:
    costs = []
    for line in data:
        costs.append(list(line.strip()))
    costs = np.array(costs).astype(int)

def a_star(costs) -> int:
    def h(point):
        m, n = point
        # approximate with straight line distance between points
        return math.sqrt((costs.shape[0]-m)**2 + (costs.shape[1]-n)**2)

    def neighbors(point):
        m, n = point
        neighboring = {(0, 1), (1, 0), (0, -1), (-1, 0)}
        return {(m+dx, n+dy) for dx, dy in neighboring
                if 0<=m+dx<costs.shape[0] and 0<=n+dy<costs.shape[1]}

    def path_sum(parents, current):
        complete = [current]
        while current in parents:
            current = parents[current]
            if current != (0, 0):
                complete.insert(0, current)
        return sum(map(lambda x: costs[x], complete))

    discovered = [(0,0)]
    parents = {}

    gScore = defaultdict(lambda : math.inf)
    gScore[(0,0)] = 0

    fScore = defaultdict(lambda : math.inf)
    fScore[(0,0)] = h((0,0))

    while discovered:
        if discovered[0] == (goal := (costs.shape[0]-1, costs.shape[1]-1)):
            return path_sum(parents, goal)

        current = heappop(discovered)
        for neighbor in neighbors(current):
            tentative = gScore[current] + costs[neighbor]
            if tentative < gScore[neighbor]:
                parents[neighbor] = current
                gScore[neighbor] = tentative
                fScore[neighbor] = tentative + h(neighbor)
                if neighbor not in discovered:
                    heappush(discovered, neighbor)
    else:
        return None
# part 1
print(a_star(costs))

# part 2
new_costs = np.zeros((costs.shape[0] * 5, costs.shape[1] * 5)).astype(int)

def increment(arr):
    arr += 1
    arr[arr > 9] = 1

xDim, yDim = costs.shape
for i in range(5):
    for j in range(5):
        copied = np.copy(costs)
        for _ in range(i+j):
            increment(copied)
        new_costs[xDim*i:xDim*(i+1), yDim*j:yDim*(j+1)] = copied

print(a_star(new_costs))