# -*- coding: utf-8 -*-
"""
Shortest paths and path lengths using A* ("A star") algorithm.

"""
__author__ =\
"""Salim Fadhley <salimfadhley@gmail.com>
Matteo Dell'Amico <matteodellamico@gmail.com>"""
#    Copyright (C) 2004-2008 by 
#    Aric Hagberg <hagberg@lanl.gov>
#    Dan Schult <dschult@colgate.edu>
#    Pieter Swart <swart@lanl.gov>
#    All rights reserved.
#    BSD license.

from heapq import heappush, heappop
from networkx import NetworkXError
import networkx as nx

__all__ = ['astar_path','astar_path_length']

def astar_path(G, source, target, heuristic=None):
    """Return a list of nodes in a shortest path between source and target using the A* ("A-star") algorithm.

    There may be more than one shortest path.  This returns only one.
    
    Parameters
    ----------
    G : NetworkX graph

    source : node
       Starting node for path

    target : node
       Ending node for path 

    heuristic : function
       A function to evaluate the estimate of the distance
       from the a node to the target.  The function takes
       two nodes arguments and must return a number.

    Raises
    ------
    NetworkXNoPath
        If no path exists between source and target.

    Examples
    --------
    >>> G=nx.path_graph(5)
    >>> print(nx.astar_path(G,0,4))
    [0, 1, 2, 3, 4]
    >>> G=nx.grid_graph(dim=[3,3])  # nodes are two-tuples (x,y)
    >>> def dist(a, b):
    ...    (x1, y1) = a
    ...    (x2, y2) = b
    ...    return ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5
    >>> print(nx.astar_path(G,(0,0),(2,2),dist))
    [(0, 0), (0, 1), (1, 1), (1, 2), (2, 2)]


    See Also
    --------
    shortest_path(), dijkstra_path()

    """
    if G.is_multigraph():
        raise NetworkXError("astar_path() not implemented for Multi(Di)Graphs")

    if heuristic is None:
        # The default heuristic is h=0 - same as Dijkstra's algorithm
        def heuristic(u,v):
            return 0
    # The queue stores priority, node, cost to reach, and parent.
    # Uses Python heapq to keep in priority order.
    queue = [(0, source, 0, None)]
    # Maps enqueued nodes to distance of discovered paths and the
    # computed heuristics to target. We avoid computing the heuristics
    # more than once and inserting the node into the queue too many times.
    enqueued = {}
    # Maps explored nodes to parent closest to the source.
    explored = {}
    
    while queue:
        # Pop the smallest item from queue.
        _, curnode, dist, parent = heappop(queue)

        if curnode == target:
            path = [curnode]
            node = parent
            while node is not None:
                path.append(node)
                node = explored[node]
            path.reverse()
            return path

        if curnode in explored:
            continue

        explored[curnode] = parent

        for neighbor, w in G[curnode].items():
            if neighbor in explored:
                continue
            ncost = dist + w.get('weight',1)
            if neighbor in enqueued:
                qcost, h = enqueued[neighbor]
                # if qcost < ncost, a longer path to neighbor remains
                # enqueued. Removing it would need to filter the whole
                # queue, it's better just to leave it there and ignore
                # it when we visit the node a second time.
                if qcost <= ncost:
                    continue
            else:
                h = heuristic(neighbor, target)
            enqueued[neighbor] = ncost, h
            heappush(queue, (ncost + h, neighbor, ncost, curnode))

    raise nx.NetworkXNoPath("Node %s not reachable from %s"%(source,target))

def astar_path_length(G, source, target, heuristic=None):
    """Return a list of nodes in a shortest path between source and target using the A* ("A-star") algorithm.

    Parameters
    ----------
    G : NetworkX graph

    source : node
       Starting node for path

    target : node
       Ending node for path 

    heuristic : function
       A function to evaluate the estimate of the distance
       from the a node to the target.  The function takes
       two nodes arguments and must return a number.

    Raises
    ------
    NetworkXNoPath
        If no path exists between source and target.

    See Also
    --------
    astar_path()

"""
    # FIXME: warn if G.weighted==False
    path=astar_path(G,source,target,heuristic)
    return sum(G[u][v].get('weight',1) for u,v in zip(path[:-1],path[1:]))
    
