import heapq
from math import sqrt


# A node within the space to search
class Node:
    # Pos is a 2-value array representing the [row,col] position of this node.
    # Parent is another instance of node.
    # Path cost is the cost to reach to reach this node.
    #   Setting path cost to zero means the total cost can be sorted such that
    #   the algorithm is a greedy search algorithm rather than A*.
    # Heuristic is the estimated benefit of using this node
    def __init__(self, pos, parent, path_cost, heuristic):
        self.pos = pos
        self.parent = parent
        self.path_cost = path_cost
        self.heuristic = heuristic
        self.total_cost = path_cost + heuristic

    # Determine if this node is the same as another
    # (If we don't do this, python will count two different objects as
    # different nodes even if they share the same location; this breaks our
    # checks for whether a node has already been visited)
    def __eq__(self, other):
        if other.pos is not None and len(other.pos) == 2:
            return other.pos[0] == self.pos[0] and other.pos[1] == self.pos[1]

    # Determine if the other node has a smaller or larger cost than another
    def __lt__(self, other):
        if other is not None:
            return self.total_cost < other.total_cost


def getNeighbors(location, cost_grid):
    # Get row count
    rows = len(cost_grid)

    # If there are no rows, return an empty list
    if rows < 1:
        return []

    # Get the columns count
    cols = len(cost_grid[0])

    # If there are no columns, return an empty list
    if cols < 1:
        return []

    # Make sure the position is valid and within the provided grid
    if len(location) != 2\
            or location[0] < 0\
            or location[0] >= rows\
            or location[1] < 0\
            or location[1] >= cols:
        return []

    # Define neighbors list
    neighbors = []

    # Add up neighbor
    if location[0] > 0 and cost_grid[location[0] - 1][location[1]] != 0:
        neighbors.append([location[0] - 1, location[1]])
    # Add down neighbor
    if location[0] < rows - 1 and cost_grid[location[0] + 1][location[1]] != 0:
        neighbors.append([location[0] + 1, location[1]])
    # Add left neighbor
    if location[1] > 0 and cost_grid[location[0]][location[1] - 1] != 0:
        neighbors.append([location[0], location[1] - 1])
    # Add right neighbor
    if location[1] < cols - 1 and cost_grid[location[0]][location[1] + 1] != 0:
        neighbors.append([location[0], location[1] + 1])

    # Return neighbors list
    return neighbors


# Expand a node into its neighbors, put them into the frontier, and mark the
# node as visited
def expandNode(node, end_pos, cost_grid, closed_list, open_list, euclidean):
    # Get possible neighbors for this node's position
    for neighbor in getNeighbors(node.pos, cost_grid):
        path_cost = cost_grid[neighbor[0]][neighbor[1]]
        if path_cost == 0:
            continue

        c = Node(neighbor, node, path_cost, calc_heuristic(neighbor, end_pos, euclidean))

        for n in open_list:
            if n.pos == neighbor:
                new_g = path_cost
                old_g = n.path_cost
                if new_g < old_g:
                    open_list.remove(n)
                    heapq.heapify(open_list)
                    heapq.heappush(open_list, c)
                    if c in closed_list:
                        closed_list.remove(c)
                        heapq.heapify(closed_list)
                break

        if c not in closed_list:
            heapq.heappush(open_list, c)


# The grid values must be separated by spaces, e.g.
# 0 0 0 0 0
# 0 1 1 1 0
# 0 1 2 1 0
# 0 0 0 0 0
# Returns a 2D list of 1s and 0s
def readGrid(filename):
    # Initialize grid to empty list
    cost_grid = []
    start_pos = [1, 1]
    end_pos = [1, 1]

    # Open the file
    with open(filename) as f:
        # Read each line of the file
        for (i, line) in enumerate(f.readlines()):
            # Split row into separate columns (by whitespace)
            cols = line.split()

            # Create an empty row for the grid
            row = []

            # Loop through the cells in this row
            for (j, cell) in enumerate(cols):
                # Set the cost for this row to 1 by default
                cost = 1

                # Check if this cell is the start cell (give it a cost of 1 for simplicity)
                if cell == 'S':
                    start_pos = [i, j]
                # Check if this cell is the goal cell (give it a cost of 1 for simplicity)
                elif cell == 'G':
                    end_pos = [i, j]
                else:
                    cost = int(cell)

                # Add this cost to the row
                row.append(cost)

            # Add the row
            cost_grid.append(row)

    # Return the grid, start, and end positions
    return cost_grid, start_pos, end_pos


# Writes a 2D list of 1s and 0s with spaces in between each character
# 1 1 1 1 1
# 1 0 0 0 1
# 1 0 0 0 1
# 1 1 1 1 1
def outputGrid(filename_str, cost_grid, start_pos, goal_pos, path):
    # Open file
    with open(filename_str, 'w') as f:
        print('Outputting grid: Goal: ' + str(goal_pos[0])
              + ', ' + str(goal_pos[1])
              + ' | Grid size: ' + str(len(cost_grid))
              + 'x' + str(len(cost_grid[0])))

        # Mark the start and goal points
        cost_grid[start_pos[0]][start_pos[1]] = 'S'
        cost_grid[goal_pos[0]][goal_pos[1]] = 'G'

        # Mark intermediate points with *
        for i, p in enumerate(path):
            if 0 < i < len(path) - 1:
                cost_grid[p[0]][p[1]] = '*'

        # Write the grid to a file
        for r, row in enumerate(cost_grid):
            for c, col in enumerate(row):
                # Don't add a ' ' at the end of a line
                if c < len(row) - 1:
                    f.write(str(col) + ' ')
                else:
                    f.write(str(col))

            # Don't add a '\n' after the last line
            if r < len(cost_grid) - 1:
                f.write("\n")


def setPath(current, path):
    # Keep track of current looping node
    node = current

    # Loop through parents to beginning
    while node is not None:
        # Add this node to the path
        path.append(node.pos)

        # Assign current node to parent
        node = node.parent

    # Flip the path around so the first element is the beginning
    path.reverse()


# Calculate the estimated cost from the provided node to the end position
# If `euclidean` is False, manhattan distance is used instead
def calc_heuristic(node_pos, end_pos, euclidean):
    if euclidean:
        return sqrt((node_pos[0] - end_pos[0])**2 + (node_pos[1] - end_pos[1])**2)
    return node_pos[0] - end_pos[0] + node_pos[1] - end_pos[1]


# Perform a search using A*
# If `euclidean` is False, manhattan distance is used instead
def informedSearch(cost_grid, start_pos, end_pos, euclidean=False):
    # Initialize search
    open_list = []
    closed_list = []

    # Add the starting node to the queue
    heapq.heappush(open_list, Node(start_pos, None, 0, calc_heuristic(start_pos, end_pos, euclidean)))

    # Loop until we reach the goal
    while len(open_list) > 0:
        # Pull a new node from open_list
        current = heapq.heappop(open_list)

        # Check if we're at the goal position
        if current.pos == end_pos:
            # Create path
            path = []
            setPath(current, path)

            # Return path
            return path

        # Add current node to the closed list
        closed_list.append(current)

        # Push traversable neighbors to `open_list`
        expandNode(current, end_pos, cost_grid, closed_list, open_list, euclidean)

    # No path was found!
    return []


def print_grid(grid):
    for r in range(len(grid)):
        for c in grid[r]:
            print(str(c), end=' ')
        print('')


def __main__():
    # Determine whether to use manhattan or euclidean distance
    getting_input = True
    euclidean = None
    while getting_input:
        euclidean = input('Use Euclidean or Manhattan distance? (E/M): ')
        if euclidean == 'E':
            euclidean = True
        elif euclidean == 'M':
            euclidean = False
        else:
            print('Please enter \"E\" or \"M\"')
            continue

        getting_input = False

    # Read the grid in, perform the search, and output the path
    print('Loading grid')
    cost_grid, start_pos, end_pos = readGrid('maze_in.txt')
    print_grid(cost_grid)
    print('Solving path')
    solved_path = informedSearch(cost_grid, start_pos, end_pos, euclidean)
    if len(solved_path) < 1:
        print('Failed to solve path!')
    else:
        print('Writing solved maze')
        outputGrid("maze_out.txt", cost_grid, start_pos, end_pos, solved_path)


__main__()