from collections import Sequence

from devito.dle.backends import (State, BasicRewriter, DevitoCustomRewriter,
                                 DevitoRewriter, DevitoSpeculativeRewriter, YaskRewriter)
from devito.exceptions import DLEException
from devito.logger import dle_warning


modes = {
    'basic': BasicRewriter,
    'advanced': DevitoRewriter,
    'speculative': DevitoSpeculativeRewriter,
    'yask': YaskRewriter
}


def transform(node, mode='basic', options=None, compiler=None):
    """
    Transform Iteration/Expression trees to generate highly optimized C code.

    :param node: The Iteration/Expression tree to be transformed, or an iterable
                 of Iteration/Expression trees.
    :param mode: Drive the tree transformation. ``mode`` is a string indicating
                 a certain optimization pipeline. The following values are accepted: ::

                     * 'noop': Do nothing.
                     * 'basic': Add instructions to avoid denormal numbers and create
                                elemental functions for rapid JIT-compilation.
                     * 'advanced': 'basic', vectorization, loop blocking.
                     * '3D-advanced': Like 'advanced', but apply 3D loop blocking
                                      if there are at least the perfectly nested
                                      parallel iteration spaces.
                     * 'speculative': Apply all of the 'advanced' transformations,
                                      plus other transformations that might increase
                                      (or possibly decrease) performance.
                     * 'yask': Optimize by offloading to the YASK optimizer. Still
                               work-in-progress; should only be used by developers.
    :param options: A dictionary with additional information to drive the DLE. The
                    following values are accepted: ::

                        * 'blockshape': A tuple representing the shape of a block created
                                        by loop blocking.
                        * 'blockinner': By default, loop blocking is not applied to the
                                        innermost dimension of an Iteration/Expression
                                        tree to maximize vectorization. Set this flag to
                                        True to override this heuristic.
                        * 'openmp': True to emit OpenMP code, False otherwise.
    :param compiler: Compiler class used to perform JIT compilation. Useful to
                     introduce compiler-specific vectorization pragmas.
    """
    assert mode is None or isinstance(mode, str)

    # Check input AST
    if isinstance(node, Sequence):
        assert all(n.is_Node for n in node)
        node = list(node)
    elif node.is_Node:
        node = [node]
    else:
        raise ValueError("Got illegal node of type %s." % type(node))

    # Check input options
    options = options or {}
    params = options.copy()
    for i in options:
        if i not in ('blockshape', 'blockinner', 'openmp'):
            dle_warning("Illegal DLE parameter '%s'" % i)
            params.pop(i)

    # Check input mode
    if mode is None or mode == 'noop':
        return State(node)
    elif mode == '3D-advanced':
        params['blockinner'] = True
        mode = 'advanced'
    elif mode not in modes:
        try:
            rewriter = DevitoCustomRewriter(node, mode, params, compiler)
            return rewriter.run()
        except DLEException:
            dle_warning("Unknown transformer mode(s) %s" % mode)
            return State(node)
    else:
        rewriter = modes[mode](node, params, compiler)
        return rewriter.run()

@@ -1213,20 +1213,22 @@
 ', but a
-pply
+ttempt
  3D loop
@@ -1240,153 +1240,8 @@
 king
-%0A                                      if there are at least the perfectly nested%0A                                      parallel iteration spaces
 .%0A  
@@ -2512,23 +2512,50 @@
 %22%22%22%0A    
-asser
+# Check input options%0A    if no
 t 
+(
 mode is 
@@ -2587,31 +2587,97 @@
 str)
-%0A%0A    # Check input AST
+):%0A        raise ValueError(%22Parameter 'mode' should be a string, not %25s.%22 %25 type(mode))%0A
 %0A   
@@ -2908,34 +2908,8 @@
 ))%0A%0A
-    # Check input options%0A
     
@@ -3137,24 +3137,61 @@
   # 
-Check input mode
+Process the Iteration/Expression tree through the DLE
 %0A   

Tweak input checks