import numpy as np
from heapq import heappush, heappop
import math


def a_star(start_node, goal_node, map, visited_map):
    start = start_node
    goal = goal_node

    # add start node to front
    # front is a list of (total estimated cost to goal, total cost from start to node, node, previous node)
    start_node_cost = 0
    start_estimated_cost_to_goal = dist2d(start, goal) + start_node_cost
    # Create first element for path, the front format [cost_to_goal, cost, pos, previous]
    front = [(start_estimated_cost_to_goal, start_node_cost, start, None)]

    # use a dictionary to remember where we came from in order to reconstruct the path later on
    came_from = {}
    movements = [(1, 0, 1.0),
                 (0, 1, 1.0),
                 (-1, 0, 1.0),
                 (0, -1, 1.0)]

    # while there are elements to investigate in our front.
    while front:
        # get smallest item and remove from front.
        element = heappop(front)

        # if this has been visited already, skip it
        total_cost, cost, pos, previous = element
        if is_visited(pos, visited_map):
            continue

        # now it has been visited, mark with cost
        visited_map = set_visited(pos, visited_map)

        # set its previous node in dict
        came_from[pos] = previous

        # if the goal has been reached, we are done!
        if pos == goal:
            break

        # check all neighbors
        for dx, dy, deltacost in movements:
            # determine new position
            new_x = pos[0] + dx
            new_y = pos[1] + dy
            new_pos = (new_x, new_y)

            # check whether new position is inside the map if not, skip node
            if not is_inside_map(new_pos, map):
                continue

            # add node to front if it was not visited before and is not an obstacle
            if (not is_visited(new_pos, visited_map)):
                potential_function_cost = map[new_pos[1]][new_pos[0]]
                new_cost = cost + deltacost + potential_function_cost
                new_total_cost_to_goal = new_cost + \
                    dist2d(new_pos, goal) + potential_function_cost

                heappush(front, (new_total_cost_to_goal, new_cost, new_pos, pos))

    # reconstruct path backwards (only if we reached the goal)
    path = []
    if pos == goal:
        while pos:
            # transform array indices to meters
            pos_m_x, pos_m_y = pos
            path.append((pos_m_x, pos_m_y))
            pos = came_from[pos]

        # reverse so that path is from start to goal.
        path.reverse()

    return path


def dist2d(point1, point2):
    # Euclidean distance between two points
    x1, y1 = point1[0:2]
    x2, y2 = point2[0:2]
    dist2 = (x1 - x2)**2 + (y1 - y2)**2
    dist = math.sqrt(dist2)
    return dist


def is_inside_map(pos, map):
    x_index, y_index = pos
    if x_index < 0 or y_index < 0 or x_index >= len(map[0]) or y_index >= len(map):
        return False
    else:
        return True


def is_visited(pos, visited_map):
    x_index, y_index = pos
    if visited_map[y_index][x_index] == 1.0:
        return True
    else:
        return False


def set_visited(pos, visited_map):
    x_index, y_index = pos
    visited_map[y_index][x_index] = 1.0
    return visited_map


def calc_risk_of_path(path, map):
    risk = 0
    # sum up all risk in path
    for i,pos in enumerate(path):
        # skip start
        if i == 0:
            continue
        risk += map[pos[1]][pos[0]]
    return risk


def main():
    # Get file in lines and convert to int
    with open("input.txt", "r", encoding="utf-8") as file:
        # strip out the new line
        puzzle_lines = [line.rstrip('\n') for line in file]

    # Generate map and visited as float
    map = np.array([list(line) for line in puzzle_lines])
    map = map.astype(np.float)
    visited_map = np.zeros((len(map), len(map[0])))
    # ([print(line) for line in map])
    # ([print(line) for line in visited_map])

    # Get path
    # set a start and an end node (in meters)
    start_node = (0, 0)
    goal_node = (len(map[0])-1, len(map)-1)

    # run A* and get the path
    path = a_star(start_node, goal_node, map, visited_map)

    # Calc risk of path
    risk = calc_risk_of_path(path, map)

    # Part
    result = risk
    # Output
    print("Result: " + str(result))


if __name__ == "__main__":
    main()
