"""
Informed search algorithms.

By Connelly Barnes, 2005.  Version 1.0.0.  Public domain.

All search algorithms return a list for a path of states from the
initial state to a goal state.  If the graph argument is True, states
are stored in a visited set, so duplicate states are not revisited.
"""

import heapq
from collections import deque

__all__ = ['State', 'path_cost', 'greedy_best_first', 'beam', 'astar']

class State:
  """
  Interface for state of environment, and state transitions.
  """
  def successors(self):
    """
    Get iterable of successor states.
    """
    raise NotImplementedError

  def cost(self, other):
    """
    Cost of state self -> other (only needed for uniform_cost_search).
    """
    raise NotImplementedError

  def is_goal(self):
    """
    True iff self is a goal state.
    """
    raise NotImplementedError

  def heuristic(self):
    """
    Estimated distance from self to goal (or other heuristic).
    """
    raise NotImplementedError


def path_cost(path):
  """
  Sum of edge costs for a given path (iterable) of states.
  """
  return sum([path[i].cost(path[i+1]) for i in xrange(len(path)-1)])


def flatten(L):
  """
  Given linked list of the form [[[[],1],2],3], return [1,2,3].
  """
  ans = []
  while len(L) > 0:
    ans.append(L[-1])
    L = L[0]
  return ans[::-1]


def greedy_best_first(initial_state, graph=False):
  """
  Greedy best-first search.
  """
  return beam(initial_state, None, graph)


def beam(initial_state, beam_width=None, graph=False):
  """
  Beam search: BFS, expand beam_width best nodes at each layer.
  """
  q = [(initial_state.heuristic(), ((),initial_state))]
  if not graph:
    while True:
      (cost, path) = heapq.heappop(q)
      state = path[-1]
      if state.is_goal():
        return flatten(path)
      for x in state.successors():
        heapq.heappush(q, (x.heuristic(), (path, x)))
      if beam_width is not None:
        q = q[:beam_width]
  else:
    visited = set()
    while True:
      (cost, path) = heapq.heappop(q)
      state = path[-1]
      if not state in visited:
        visited.add(state)
        if state.is_goal():
          return flatten(path)
        for x in state.successors():
          if x not in visited:
            heapq.heappush(q, (x.heuristic(), (path, x)))
      if beam_width is not None:
        q = q[:beam_width]


def astar(initial_state, graph=False):
  """
  A* search.

  For A* to be optimal in tree search mode, the state's heuristic()
  must be <= the cost of the minimal path to the goal node.  For A* to
  be optimal in graph search mode, for all states a and successors b:

    0 <= a.heuristic() - b.heuristic() <= a.cost(b)

  It suffices for vertices to be in a normed vector space, with
  a.cost(b) == norm(a - b) and a.heuristic() == norm(a - goal).
  """
  q = [(initial_state.heuristic(), 0, ((),initial_state))]
  if not graph:
    while True:
      (total, precost, path) = heapq.heappop(q)
      state = path[-1]
      if state.is_goal():
        return flatten(path)
      for x in state.successors():
        precost2 = precost + state.cost(x)
        heapq.heappush(q, (precost2 + x.heuristic(), precost2,
                           (path, x)))
  else:
    visited = set()
    while True:
      (total, precost, path) = heapq.heappop(q)
      state = path[-1]
      if not state in visited:
        visited.add(state)
        if state.is_goal():
          return flatten(path)
        for x in state.successors():
          if x not in visited:
            precost2 = precost + state.cost(x)
            heapq.heappush(q, (precost2 + x.heuristic(), precost2,
                               (path, x)))