from devito.symbolics.queries import (q_indexed, q_function, q_terminal, q_leaf, q_xop,
                                      q_scalar)
from devito.tools import as_tuple

__all__ = ['retrieve_indexed', 'retrieve_functions', 'retrieve_function_carriers',
           'retrieve_terminals', 'retrieve_xops', 'retrieve_scalars', 'search']


class Search(object):

    class Set(set):

        @staticmethod
        def wrap(obj):
            return {obj}

    class List(list):

        @staticmethod
        def wrap(obj):
            return [obj]

        def update(self, obj):
            return self.extend(obj)

    modes = {
        'unique': Set,
        'all': List
    }

    def __init__(self, query, mode, deep=False):
        """
        Search objects in an expression. This is much quicker than the more
        general SymPy's find.

        Parameters
        ----------
        query
            Any query from :mod:`queries`.
        mode : str
            Either 'unique' or 'all' (catch all instances).
        deep : bool, optional
            If True, propagate the search within an Indexed's indices. Defaults to False.
        """
        self.query = query
        self.collection = self.modes[mode]
        self.deep = deep

    def _next(self, expr):
        if self.deep is True and expr.is_Indexed:
            return expr.indices
        elif q_leaf(expr):
            return ()
        else:
            return expr.args

    def dfs(self, expr):
        """
        Perform a DFS search.

        Parameters
        ----------
        expr : expr-like
            The searched expression.
        """
        found = self.collection()
        for a in self._next(expr):
            found.update(self.dfs(a))
        if self.query(expr):
            found.update(self.collection.wrap(expr))
        return found

    def bfs(self, expr):
        """
        Perform a BFS search.

        Parameters
        ----------
        expr : expr-like
            The searched expression.
        """
        found = self.collection()
        if self.query(expr):
            found.update(self.collection.wrap(expr))
        for a in self._next(expr):
            found.update(self.bfs(a))
        return found

    def bfs_first_hit(self, expr):
        """
        Perform a BFS search, returning immediately when a node matches the query.

        Parameters
        ----------
        expr : expr-like
            The searched expression.
        """
        found = self.collection()
        if self.query(expr):
            found.update(self.collection.wrap(expr))
            return found
        for a in self._next(expr):
            found.update(self.bfs_first_hit(a))
        return found


def search(exprs, query, mode='unique', visit='dfs', deep=False):
    """Interface to Search."""

    assert mode in Search.modes, "Unknown mode"
    assert visit in ['dfs', 'bfs', 'bfs_first_hit'], "Unknown visit type"

    searcher = Search(query, mode, deep)

    found = Search.modes[mode]()
    for e in as_tuple(exprs):
        if visit == 'dfs':
            found.update(searcher.dfs(e))
        elif visit == 'bfs':
            found.update(searcher.bfs(e))
        else:
            found.update(searcher.bfs_first_hit(e))

    return found


# Shorthands


def retrieve_indexed(exprs, mode='all', deep=False):
    """Shorthand to retrieve the Indexeds in ``exprs``."""
    return search(exprs, q_indexed, mode, 'dfs', deep)


def retrieve_functions(exprs, mode='all'):
    """Shorthand to retrieve the DiscreteFunctions in ``exprs``."""
    return search(exprs, q_function, mode, 'dfs')


def retrieve_scalars(exprs, mode='all'):
    """Shorthand to retrieve the Scalar in ``exprs``."""
    return search(exprs, q_scalar, mode, 'dfs')


def retrieve_function_carriers(exprs, mode='all'):
    """
    Shorthand to retrieve the DiscreteFunction carriers in ``exprs``. An
    object carries a DiscreteFunction if any of the following conditions are met: ::

        * it is itself a DiscreteFunction, OR
        * it is an Indexed, which internally has a pointer to a DiscreteFunction.
    """
    query = lambda i: q_function(i) or q_indexed(i)
    retval = search(exprs, query, mode, 'dfs')
    # Filter off Indexeds not carrying a DiscreteFunction
    for i in list(retval):
        try:
            i.function
        except AttributeError:
            retval.remove(i)
    return retval


def retrieve_terminals(exprs, mode='all', deep=False):
    """Shorthand to retrieve Indexeds and Symbols within ``exprs``."""
    return search(exprs, q_terminal, mode, 'dfs', deep)


def retrieve_xops(exprs):
    """Shorthand to retrieve the arithmetic operations within ``exprs``."""
    return search(exprs, q_xop, 'all', 'dfs')

@@ -1825,32 +1825,510 @@
   return found%0A%0A
+    def dfs_first_hit(self, expr):%0A        %22%22%22%0A        Perform a DFS search, returning immediately when a node matches the query.%0A%0A        Parameters%0A        ----------%0A        expr : expr-like%0A            The searched expression.%0A        %22%22%22%0A        found = self.collection()%0A        for a in self._next(expr):%0A            found.update(self.dfs_first_hit(a))%0A        if not found and self.query(expr):%0A            found.update(self.collection.wrap(expr))%0A        return found%0A%0A
     def bfs(self
@@ -3341,82 +3341,8 @@
 ode%22
-%0A    assert visit in %5B'dfs', 'bfs', 'bfs_first_hit'%5D, %22Unknown visit type%22
 %0A%0A  
@@ -3509,24 +3509,115 @@
 her.dfs(e))%0A
+        elif visit == 'dfs_first_hit':%0A            found.update(searcher.dfs_first_hit(e))%0A
         elif
@@ -3677,34 +3677,59 @@
 s(e))%0A        el
-se
+if visit == %22bfs_first_hit%22
 :%0A            fo
@@ -3765,16 +3765,94 @@
 _hit(e))
+%0A        else:%0A            raise ValueError(%22Unknown visit type %60%25s%60%22 %25 visit)
 %0A%0A    re

Add dfs_first_hit search