import heapq

"""
Explanation: Weighted A$^*$ Search
The module heapq provides 
priority queues 
that are implemented as heaps.
Technically, these heaps are just lists.  In order to use them as priority queues, the entries of these lists will be pairs of the form $(p, o)$, where $p$ is the priority of the object $o$.  Usually, the priorities are numbers 
and, contra-intuitively, high priorities correspond to <b>small</b> numbers, that is $(p_1, o_1)$ has a higher priority than $(p_2, o_2)$ iff $p_1 < p_2$.
We need only two functions from the module heapq:
- Given a heap $H$, the function $\texttt{heapq.heappop}(H)$ removes the pair
  from H that has the highest priority.  This pair is also returned.
- Given a heap $H$, the function $\texttt{heapq.heappush}\bigl(H, (p, o)\bigr)$ 
  pushes the pair $(p, o)$ onto the heap $H$.  This method does not return a 
  value.  Instead, the heap $H$ is changed in place.
End of explanation
"""


def search(start, goal, next_states, heuristic):
    Visited   = set()
    PrioQueue = [ (2*heuristic(start, goal), [start]) ]
    while PrioQueue:
        _, Path = heapq.heappop(PrioQueue)
        state   = Path[-1]
        if state in Visited:
            continue
        if state == goal:
            return Path
        for ns in next_states(state):           
            if ns not in Visited:
                prio = 2*heuristic(ns, goal) + len(Path) + 1
                heapq.heappush(PrioQueue, (prio, Path + [ns]))
        Visited.add(state)

%run Sliding-Puzzle.ipynb

%load_ext memory_profiler

%%time
%memit Path = search(start, goal, next_states, manhattan)
print(len(Path)-1)

animation(Path)

%%time
Path = search(start2, goal2, next_states, manhattan)
print(len(Path)-1)

animation(Path)

"""
Explanation: The function search takes three arguments to solve a search problem:
- start is the start state of the search problem,
- goal is the goal state, and
- next_states is a function with signature $\texttt{next_states}:Q \rightarrow 2^Q$, where $Q$ is the set of states.
  For every state $s \in Q$, $\texttt{next_states}(s)$ is the set of states that can be reached from $s$ in one step.
- heuristic is a function that takes two states as arguments.  It returns an estimate of the 
  length of the shortest path between these states.
If successful, search returns a path from start to goal that is a solution of the search problem
$$ \langle Q, \texttt{next_states}, \texttt{start}, \texttt{goal} \rangle. $$
The variable PrioQueue that is used in the implementation contains pairs of the form
$$ \bigl(\texttt{len}(\texttt{Path}) + \texttt{heuristic}(\texttt{state},\; \texttt{goal}), \texttt{Path}\bigr), $$
where Path is a path from start to state and $\texttt{heuristic}(\texttt{state}, \texttt{goal})$
is an estimate of the distance from state to goal.  The idea is to always extend the most promising Path, i.e. to extend the Path whose completed version would be shortest.
End of explanation
"""


