#
# Tree-search Python module
#

# Copyright [2006] [alex at delandgraaf dot com]
# Licensed under the LGPL

# Usage:
# import treesearch
#
# search = treesearch.TreeSearch()
# search.dfs() # Depth-first search
# search.bfs() # Breadth-first search

# Defaults to an 'a' starting element, child_function returns the next
# character in the alphabet and eval_function checks if the element is 'z'

# Useful functions:
# search = treesearch.TreeSearch(top_element) # Initialize using given element
# search.set_eval_function(my_eval_function)
# search.set_child_function(my_child_function)

# Optionally, you can give the max number of evaluations to dfs() and bfs()

# Ideally, you should only have to implement your own
# evaluation and child-functions, give your first element, and use
# the search-function you prefer :)

# During implementation, use search.debug = 1 to output useful stuff

###

# Implementation details:

# A Treesearch can be implemented in a number of ways.
# The easiest in my opinion is the use of a simple stack:
# Push the initial state on to the stack.
# Then, for every top element on the stack evaluate the element.
# If it is a valid solution, stop.
# If not, remove the top element of the stack and add all it's children
# to the top of the stack (or, with bfs, shove all children onto the bottom)
# Continue until the stack is empty, in which case there is no answer.

# Yes, I know that only DFS is a true stack, as you can't add elements
# to the bottom of one as with BFS. Go back and play with your lego, kiddies.

class TreeSearch:
    """ Generic TreeSearch class, containing all that is required
    for building a simple search tool

    set_eval_function(eval_function)
      eval_function should, given an element:
       return 0 on unsuccessful answer, 1 on answer

    set_child_function(child_function)
      child_function should, given an element:
       return a list of children, or an empty list if there are none
    """

    stack = []
    eval_function = lambda x: 0
    child_function = lambda x: 0
    mode = "dfs"
    debug = 0
    max_evaluations = -1
    nr_elements_evaluated = 0

    def __init__(self, top_element = "a"):
        self.stack.append(top_element)
        
        self.eval_function = self.default_eval_function
        self.child_function = self.default_child_function

    def set_eval_function(self, x):
        self.eval_function = x

    def set_child_function(self, x):
        self.child_function = x
    
    def default_eval_function(self, element):
        if element == "z":
            return 1
        return 0

    def default_child_function(self, element):
        children = []
        if ord(element[0]) <= ord('y'):
            children.append(chr(ord(element[0]) + 1))
        return children

    def dfs(self, max_evaluations = -1):
        self.max_evaluations = max_evaluations
        self.mode = "dfs"
        self.nr_elements_evaluated = 0
        return self.search(self.stack)

    def bfs(self, max_evaluations = -1):
        self.max_evaluations = max_evaluations        
        self.mode = "bfs"
        self.nr_elements_evaluated = 0        
        return self.search(self.stack)

    def pdebug(self, string):
        if self.debug == 1:
            print string

    def search(self, stack):
        """ Peform the search recursively:
        - if no answer can be found (or max_evaluations has been reached):
            return -1
        - if the answer is found,
            return the succesful element
               (might be useful to add some bookkeeping here)
        """

        self.pdebug("Searching Tree with stack: " + str(stack))
        self.pdebug("Evaluations: " + str(self.nr_elements_evaluated) + ", max evaluations: " + str(self.max_evaluations))

        if len(stack) == 0:
            self.pdebug("Stack empty, didn't find a solution")
            return -1
        
        if self.max_evaluations != -1 and self.nr_elements_evaluated >= self.max_evaluations:
            self.pdebug("Maximum number of element evaluations reached, aborting")
            return -1
        

        self.nr_elements_evaluated += 1
        
        top = stack.pop()        
        if self.eval_function(top) == 1:
            self.pdebug("Found a solution, returning")
            return top
        
        
        children = self.child_function(top)
        
        self.pdebug("Generated children: " + str(children))
        if self.mode == "dfs":
            for i in children:
                stack.append(i)
                
        elif self.mode == "bfs":
            stack = children + stack

        self.pdebug("New stack: " + str(stack))

        return self.search(stack)