from hashlib import sha1
from itertools import chain
from os import path

import cgen as c
import numpy as np

from devito.compiler import (get_compiler_from_env, get_tmp_dir,
                             jit_compile_and_load)
from devito.expression import Expression
from devito.interfaces import SymbolicData
from devito.iteration import Iteration
from devito.logger import error
from devito.tools import filter_ordered

__all__ = ['StencilKernel']


class StencilKernel(object):
    """Code generation class, alternative to Propagator"""

    def __init__(self, stencils, name="Kernel", compiler=None):
        # Default attributes required for compilation
        self.name = name
        self.compiler = compiler or get_compiler_from_env()
        self._lib = None
        self._cfunction = None
        # Ensure we always deal with Expression lists
        stencils = stencils if isinstance(stencils, list) else [stencils]
        self.expressions = [Expression(s) for s in stencils]

        # Lower all expressions to "indexed" API
        for e in self.expressions:
            e.indexify()

        # Wrap expressions with Iterations according to dimensions
        for i, expr in enumerate(self.expressions):
            newexpr = expr
            offsets = newexpr.index_offsets
            for d in reversed(expr.dimensions):
                newexpr = Iteration(newexpr, d, d.size, offsets=offsets[d])
            self.expressions[i] = newexpr

        # TODO: Merge Iterations iff outermost variables agree

    def __call__(self, *args, **kwargs):
        self.apply(*args, **kwargs)

    def apply(self, *args, **kwargs):
        """Apply defined stencil kernel to a set of data objects"""
        if len(args) <= 0:
            args = self.signature
        args = [a.data if isinstance(a, SymbolicData) else a for a in args]
        # Check shape of argument data
        for arg, v in zip(args, self.signature):
            if not isinstance(arg, np.ndarray):
                raise TypeError('No array data found for argument %s' % v.name)
            if arg.shape != v.shape:
                error('Expected argument %s with shape %s, but got shape %s'
                      % (v.name, v.shape, arg))
                raise ValueError('Argument with wrong shape')
        # Invoke kernel function with args
        self.cfunction(*args)

    @property
    def signature(self):
        """List of data objects that define the kernel signature

        :returns: List of unique data objects required by the kernel
        """
        signatures = [e.signature for e in self.expressions]
        return filter_ordered(chain(*signatures))

    @property
    def ccode(self):
        """Returns the C code generated by this kernel.

        This function generates the internal code block from Iteration
        and Expression objects, and adds the necessary template code
        around it.
        """
        header_vars = [c.Pointer(c.POD(v.dtype, '%s_vec' % v.name))
                       for v in self.signature]
        header = c.Extern("C", c.FunctionDeclaration(
            c.Value('int', self.name), header_vars))
        cast_shapes = [(v, ''.join(['[%d]' % d for d in v.shape[1:]]))
                       for v in self.signature]
        casts = [c.Initializer(c.POD(v.dtype, '(*%s)%s' % (v.name, shape)),
                               '(%s (*)%s) %s' % (c.dtype_to_ctype(v.dtype),
                                                  shape, '%s_vec' % v.name))
                 for v, shape in cast_shapes]
        body = [e.ccode for e in self.expressions]
        ret = [c.Statement("return 0")]
        return c.FunctionBody(header, c.Block(casts + body + ret))

    @property
    def basename(self):
        """Generate the file basename path for auto-generated files

        The basename is generated from the hash string of the kernel,
        which is base on the final expressions and iteration symbols.

        :returns: The basename path as a string
        """
        expr_string = "\n".join([str(e) for e in self.expressions])
        hash_key = sha1(expr_string.encode()).hexdigest()

        return path.join(get_tmp_dir(), hash_key)

    @property
    def cfunction(self):
        """Returns the JIT-compiled C function as a ctypes.FuncPtr object

        Note that this invokes the JIT compilation toolchain with the
        compiler class derived in the constructor

        :returns: The generated C function
        """
        if self._lib is None:
            self._lib = jit_compile_and_load(self.ccode, self.basename,
                                             self.compiler)
        if self._cfunction is None:
            self._cfunction = getattr(self._lib, self.name)
            self._cfunction.argtypes = self.argtypes

        return self._cfunction

    @property
    def argtypes(self):
        """Create argument types for defining function signatures via ctypes

        :returns: A list of ctypes of the matrix parameters and scalar parameters
        """
        return [np.ctypeslib.ndpointer(dtype=v.dtype, flags='C')
                for v in self.signature]

@@ -530,16 +530,593 @@
 opagator
+%0A%0A    :param stencils: SymPy equation or list of equations that define the%0A                     stencil used to create the kernel of this Operator.%0A    :param name: Name of the kernel function - defaults to %22Kernel%22%0A    :param subs: Dict or list of dicts containing the SymPy symbol%0A                 substitutions for each stencil respectively.%0A    :param compiler: Compiler class used to perform JIT compilation.%0A                     If not provided, the compiler will be inferred from the%0A                     environment variable DEVITO_ARCH, or default to GNUCompiler.%0A    
 %22%22%22%0A%0A   
@@ -1159,16 +1159,27 @@
 Kernel%22,
+ subs=None,
  compile
@@ -1678,24 +1678,173 @@
 indexify()%0A%0A
+        # Apply supplied substitutions%0A        if subs is not None:%0A            for expr in self.expressions:%0A                expr.substitute(subs)%0A%0A
         # Wr

Add global symbol substitution to code-generation stage