#!/usr/bin/env python
# coding: utf-8

# # Informed Search

# Also known as "heuristic" search, because the search is informed by an
# estimate of the total path cost through each node, and the next
# unexpanded node with the lowest estimated cost is expanded next. 
# 
#     At some intermediate node, the 
#       estimated cost of the solution path =
#           the sum of the step costs so far from the start node to this node
#              +
#           an estimate of the sum of the remaining step costs to a goal
# 
# Let's label these as
# 
#    * $f(n) =$ estimated cost of the solution path through node $n$
#    * $g(n) =$ the sum of the step costs so far from the start node to this node
#    * $h(n) =$ an estimate of the sum of the remaining step costs to a goal
# 
# *heuristic function*: $h(n) =$ estimated cost of the cheapest path from state at node $n$ to a goal state.
# 
# <img src='https://www.cs.colostate.edu/~anderson/cs440/notebooks/astarfgh.png'>
# 
# Should we explore under Node a or b?
# 
# <img src='https://www.cs.colostate.edu/~anderson/cs440/notebooks/astarfgh2.png'>

# # A* algorithm

# ## Non-recursive

# So, now you know enough python to try to implement A\*, at least a non-recursive form.  Start with your graph search algorithm from Assignment 1.  Modify it so that the next node selected is based on its `f` value.
# 
# For a given problem, define `start_state`, `actions_f`, `take_action_f`, `goal_test_f`, and a heuristic function `heuristic_f`.  `actions_f` must return valid actions paired with the single step cost, and `take_action_f` must return the pair containing the new state and the cost of the single step given by the action.  We can use the `Node` class to hold instances of nodes.  However, since this is not a recursive algorithm, `Node` must be extended to include the node's parent node, to be able to generate the solution path once the search finds the goal.
# 
# Now the A* algorithm can be written as follows
#   * Initialize `expanded` to be an empty dictionary
#   * Initialize `un_expanded` to be a list containing the start_state node.  Its `h` value is calculated using `heuristic_f`, its `g` value is 0, and its `f` value is `g+h`.
#   * If `start_state` is the `goal_state`, return the list containing just `start_state` and its `f` value to show the cost of the solution path.
#   * Repeat the following steps while `un_expanded` is not empty:
#     * Pop from the front of `un_expanded` to get the best (lowest f value) node to expand.
#     * Generate the `children` of this `node`.
#     * Update the `g` value of each child by adding the action's single step cost to this node's `g` value.
#     * Calculate `heuristic_f` of each child.
#     * Set `f = g + h` of each child.
#     * Add the node to the `expanded` dictionary, indexed by its state.
#     * Remove from `children` any nodes that are already either in `expanded` or `un_expanded`, unless the node in `children` has a lower f value.
#     * If `goal_state` is in `children`:
#       * Build the solution path as a list starting with `goal_state`.  
#       * Use the parent stored with each node in the `expanded` dictionary to construct the path.
#       * Reverse the solution path list and return it.
#     * Insert the modified `children` list into the `un_expanded`   list and ** sort by `f` values.**

# ## Recursive

# Our authors provide the Recursive Best-First Search algorithm, which
# is A\* in a recursive, iterative-deepening form, where depth is now
# given by the $f$ value.  Other differences from just
# iterative-deepening A\* are:
#   - depth-limit determined by $f$ value of best alternative to node being explored, so will stop when alternative at the node's level looks better;
#   - $f$ value of a node is replaced by best $f$ value of its children, so any future decision to try expanding this node again is more informed.
# 
# It is a bit difficult to translate their pseudo-code into python.  Here is my version.  Let's step through it.

# In[32]:


get_ipython().run_cell_magic('writefile', 'a_star_search.py', '# Recursive Best First Search (Figure 3.26, Russell and Norvig)\n#  Recursive Iterative Deepening form of A*, where depth is replaced by f(n)\n\nclass Node:\n\n    def __init__(self, state, f=0, g=0, h=0):\n        self.state = state\n        self.f = f\n        self.g = g\n        self.h = h\n\n    def __repr__(self):\n        return f\'Node({self.state}, f={self.f}, g={self.g}, h={self.h})\'\n\ndef a_star_search(start_state, actions_f, take_action_f, goal_test_f, heuristic_f):\n    h = heuristic_f(start_state)\n    start_node = Node(state=start_state, f=0 + h, g=0, h=h)\n    return a_star_search_helper(start_node, actions_f, take_action_f, \n                                goal_test_f, heuristic_f, float(\'inf\'))\n\ndef a_star_search_helper(parent_node, actions_f, take_action_f, \n                         goal_test_f, heuristic_f, f_max):\n\n    if goal_test_f(parent_node.state):\n        return ([parent_node.state], parent_node.g)\n    \n    ## Construct list of children nodes with f, g, and h values\n    actions = actions_f(parent_node.state)\n    if not actions:\n        return (\'failure\', float(\'inf\'))\n    \n    children = []\n    for action in actions:\n        (child_state, step_cost) = take_action_f(parent_node.state, action)\n        h = heuristic_f(child_state)\n        g = parent_node.g + step_cost\n        f = max(h + g, parent_node.f)\n        child_node = Node(state=child_state, f=f, g=g, h=h)\n        children.append(child_node)\n        \n    while True:\n        # find best child\n        children.sort(key = lambda n: n.f) # sort by f value\n        best_child = children[0]\n        if best_child.f > f_max:\n            return (\'failure\', best_child.f)\n        # next lowest f value\n        alternative_f = children[1].f if len(children) > 1 else float(\'inf\')\n        # expand best child, reassign its f value to be returned value\n        result, best_child.f = a_star_search_helper(best_child, actions_f,\n                                                    take_action_f, goal_test_f,\n                                                    heuristic_f,\n                                                    min(f_max,alternative_f))\n        if result != \'failure\':                    #        g\n            result.insert(0, parent_node.state)    #       / \n            return (result, best_child.f)          #      d\n                                                   #     / \\ \nif __name__ == "__main__":                         #    b   h   \n                                                   #   / \\   \n    successors = {\'a\': [\'b\',\'c\'],                  #  a   e  \n                  \'b\': [\'d\',\'e\'],                  #   \\         \n                  \'c\': [\'f\'],                      #    c   i\n                  \'d\': [\'g\', \'h\'],                 #     \\ / \n                  \'f\': [\'i\',\'j\']}                  #      f  \n                                                   #      \\\n    def actions_f(state):                          #       j \n        try:\n            ## step cost of each action is 1\n            return [(succ, 1) for succ in successors[state]]\n        except KeyError:\n            return []\n\n    def take_action_f(state, action):\n        return action\n\n    def goal_test_f(state):\n        return state == goal\n\n    def h1(state):\n        return 0\n\n    start = \'a\'\n    goal = \'h\'\n    result = a_star_search(start, actions_f, take_action_f, goal_test_f, h1)\n\n    print(f\'Path from a to h is {result[0]} for a cost of {result[1]}\')\n')


# Running this shows

# In[33]:


run a_star_search.py


# In[34]:


actions_f('a')
valid_ones = actions_f('a')
valid_ones


# In[35]:


take_action_f('a', valid_ones[0])


# Actually, there is in error in this code.  Try using it to search for a goal that does not exist!