import string

import sys, os; sys.path.append(os.path.join(os.path.dirname(__file__), "../lib"))
import astar

elevation = {c:i for i, c in enumerate(string.ascii_lowercase)}
elevation['E'] = elevation['z']
elevation['S'] = elevation['a']

data = open("input").read().splitlines(False)
#print(*data, sep="\n")

extra = []
for y, row in enumerate(data):
    for x, c in enumerate(row):
        if c == 'S':
            start = (x,y)
        elif c == 'E':
            goal = (x,y)
        elif c == 'a':
            # for part 2
            extra.append((x,y))

def get(x, y):
    return elevation[data[y][x]]

def solve(part):
    def neighbors(src):
        x,y = src
        here = get(x,y)
        n = []
        def push(x,y):
            if 0 <= y < len(data):
                if 0 <= x < len(data[y]):
                    if get(x,y) <= here+1:
                        n.append((1, (x,y)))
        push(x-1, y)
        push(x+1,y)
        push(x, y-1)
        push(x, y+1)
        return n

    def heuristic(n):
        return abs(goal[0] - n[0]) + abs(goal[1] - n[1])

    s = [start]
    if part == 2:
        s += extra

    print(part, "=", astar.search(s, goal, neighbors, heuristic))

solve(1)
solve(2)