
"""
Python bindings for Halide
"""

from __future__ import print_function
from __future__ import division


from .cHalide import *               # The direct wrapper of the C++ code generated by SWIG
import numpy
try:
    import Image as PIL
except ImportError:
    import PIL.Image as PIL

import os
import pkgutil
try:
    import cStringIO
except ImportError:
    import io
import tempfile
import sys

# ----------------------------------------------------------------------------------------------------------

if sys.version < '3':
    # python 2
    integer_types = (int, long,)
    str_types = (str, unicode)
else:
    # python 3
    integer_types = (int,)
    str_types = (str,)

# ----------------------------------------------------------------------------------------------------------

__version__ = '0.3.0'

exit_on_signal()                   # Install C++ debugging traceback

# ----------------------------------------------------------------------------------------------------------
# Types (classes/functions are used to replace the "constructors" due to working around SWIG)
# ----------------------------------------------------------------------------------------------------------

ExprType = Expr
ParamTypes = (Param_int8, Param_int16, Param_int32, Param_uint8, Param_uint16, Param_uint32, Param_float32, Param_float64)
RDomType = RDom
ImageTypes = (Image_int8, Image_int16, Image_int32, Image_uint8, Image_uint16, Image_uint32, Image_float32, Image_float64)
ImageParamType = ImageParam
BufferType = Buffer
FuncType = Func
VarType = Var
RVarType = RVar
RDomType = RDom
TypeType = Type
FuncRefExprType = FuncRefExpr
FuncRefVarType = FuncRefVar
VarOrRVarType = VarOrRVar
ArgumentType = Argument
StageType = Stage

GPU_Default = 0
GPU_CUDA = 1
GPU_OpenCL = 2
GPU_GLSL = 3

# ----------------------------------------------------
# Expr
# ----------------------------------------------------

def wrap(*a):
#    print(a)
    if len(a) == 1:
        if isinstance(a[0], ParamTypes):
            return cast_to_expr(a[0])
        elif isinstance(a[0], ImageTypes):# or isinstance(a[0], DynImageType):
            return a[0] #ExprType(to_dynimage(a[0]))
        elif isinstance(a[0], integer_types):
            return cast_to_expr(a[0])
        else:
            return cast_to_expr(a[0])
            #return expr_from_tuple(*(wrap(x) for x in a[0]))
    return ExprType(*a)

pow0 = pow

_lerp0 = lerp
lerp =  lambda zero_value, one_value, weight: _lerp0(wrap(zero_value), wrap(one_value), weight)
"""
Linear interpolate between the two values according to a weight.
  param zero_val The result when weight is 0
  param one_val The result when weight is 1
  param weight The interpolation amount

"""

for BaseT in (ExprType, FuncRefExpr, FuncRefVar, VarType, RDomType, RVarType, FuncType) + ParamTypes:
    BaseT.__add__ = lambda x, y: add(wrap(x), wrap(y))
    BaseT.__sub__ = lambda x, y: sub(wrap(x), wrap(y))
    BaseT.__mul__ = lambda x, y: mul(wrap(x), wrap(y))
    BaseT.__div__ = lambda x, y: div(wrap(x), wrap(y))
    BaseT.__mod__ = lambda x, y: mod(wrap(x), wrap(y))
    BaseT.__pow__ = lambda x, y: pow0(wrap(x), wrap(y))
    BaseT.__and__  = lambda x, y: and_op(wrap(x), wrap(y))
    BaseT.__or__  = lambda x, y: or_op(wrap(x), wrap(y))

    BaseT.__radd__ = lambda y, x: add(wrap(x), wrap(y))
    BaseT.__rsub__ = lambda y, x: sub(wrap(x), wrap(y))
    BaseT.__rmul__ = lambda y, x: mul(wrap(x), wrap(y))
    BaseT.__rdiv__ = lambda y, x: div(wrap(x), wrap(y))
    BaseT.__rmod__ = lambda y, x: mod(wrap(x), wrap(y))
    BaseT.__rpow__ = lambda y, x: pow(wrap(x), wrap(y))
    BaseT.__rand__  = lambda y, x: and_op(wrap(x), wrap(y))
    BaseT.__ror__  = lambda y, x: or_op(wrap(x), wrap(y))

    BaseT.__neg__ = lambda x: neg(wrap(x))
    BaseT.__invert__  = lambda x: invert(wrap(x))

    BaseT.__lt__  = lambda x, y: lt(wrap(x), wrap(y))
    BaseT.__le__  = lambda x, y: le(wrap(x), wrap(y))
    BaseT.__eq__  = lambda x, y: eq(wrap(x), wrap(y))
    BaseT.__ne__  = lambda x, y: ne(wrap(x), wrap(y))
    BaseT.__gt__  = lambda x, y: gt(wrap(x), wrap(y))
    BaseT.__ge__  = lambda x, y: ge(wrap(x), wrap(y))

    BaseT.__iadd__ = lambda x, y: iadd(wrap(x), wrap(y))
    BaseT.__isub__ = lambda x, y: isub(wrap(x), wrap(y))
    BaseT.__imul__ = lambda x, y: imul(wrap(x), wrap(y))
    BaseT.__idiv__ = lambda x, y: idiv(wrap(x), wrap(y))

# ----------------------------------------------------
# Var, Func
# ----------------------------------------------------

class Var(object):
    """
    A Halide variable, to be used when defining functions. It is just
    a name, and can be reused in places where no name conflict will
    occur. It can be used in the left-hand-side of a function
    definition, or as an Expr. As an Expr, it always has type Int(32).

    Constructors::

      Var()      -- Construct Var with an automatically-generated unique name
      Var(name)  -- Construct Var with the given string name.
    """
    def __new__(cls, *args):
        return VarType(*args)

    def name(self):
        """
        Get the name of a Var.
        """

    def same_as(self, other):
        """
        Test if two Vars are the same.
        """

    def implicit(self, n):
        """
        Construct implicit Var from int n.
        """

for C in [VarType, FuncRefExpr, FuncRefVar]:
    C.__add__ = lambda x, y: add(wrap(x), wrap(y))
    C.__sub__ = lambda x, y: sub(wrap(x), wrap(y))
    C.__mul__ = lambda x, y: mul(wrap(x), wrap(y))
    C.__div__ = lambda x, y: div(wrap(x), wrap(y))
    C.__mod__ = lambda x, y: mod(wrap(x), wrap(y))
    C.__pow__ = lambda x, y: pow(wrap(x), wrap(y))
    C.__and__  = lambda x, y: and_op(wrap(x), wrap(y))
    C.__or__  = lambda x, y: or_op(wrap(x), wrap(y))

    C.__neg__ = lambda x: neg(wrap(x))
    C.__invert__  = lambda x: invert(wrap(x))

    C.__lt__  = lambda x, y: lt(wrap(x), wrap(y))
    C.__le__  = lambda x, y: le(wrap(x), wrap(y))
    C.__eq__  = lambda x, y: eq(wrap(x), wrap(y))
    C.__ne__  = lambda x, y: ne(wrap(x), wrap(y))
    C.__gt__  = lambda x, y: gt(wrap(x), wrap(y))
    C.__ge__  = lambda x, y: ge(wrap(x), wrap(y))

# ----------------------------------------------------
# Func
# ----------------------------------------------------

def raise_error(e):
    raise e

def wrap_func_arg(a, useInt=True):
#    print(a)
    if isinstance(a, VarOrRVarType):
        return a
    elif isinstance(a, RVarType):
        return VarOrRVar(a)
    elif isinstance(a, VarType):
        return VarOrRVar(a)
    elif useInt and isinstance(a, integer_types):
        return a
    return wrap(a)

#_reset0 = Func.reset     # Reset schedule has been removed (could be reimplemented in C++ layer as a walk over all Funcs)
_generic_getitem_var_or_expr = lambda x, key: call(x, *wrap_func_args(key)) if isinstance(key,tuple) else call(x, wrap(key) if not isinstance(key,VarType) else key)
_generic_getitem_expr = lambda x, key: call(x, *[wrap(y) for y in key]) if isinstance(key,tuple) else call(x, wrap(key))
_generic_set = lambda x, y: set(x, wrap(y))
#_realize = Func.realize

FuncType.__call__ = lambda self, *L: raise_error(ValueError('used f(x, y) to refer to a Func -- proper syntax is f[x, y]'))

def wrap_func_args(args):
    if all(isinstance(x, VarType) for x in args):
        return args
    return [wrap(y) for y in args]

def wrap_gpu_args(args):
    return [wrap_func_arg(y, False) for y in args]

def wrap_gpu_args_int(args):
    return [wrap_func_arg(y) for y in args]

FuncType.__setitem__ = lambda x, key, value: set(call(x, *wrap_func_args(key)), wrap(value)) if isinstance(key,tuple) else set(call(x, wrap(key) if not isinstance(key,VarType) else key), wrap(value))
FuncType.__getitem__ = _generic_getitem_var_or_expr
FuncType.set = _generic_set

_allow_race_conditions0 = FuncType.allow_race_conditions
_gpu0 = FuncType.gpu
_gpu_blocks0 = FuncType.gpu_blocks
_gpu_single_thread0 = FuncType.gpu_single_thread
_gpu_threads0 = FuncType.gpu_threads
_gpu_tile0 = FuncType.gpu_tile
_serial0 = FuncType.serial
_split0 = FuncType.split
_fuse0 = FuncType.fuse
_tile0 = FuncType.tile
_reorder0 = FuncType.reorder
_bound0 = FuncType.bound
_parallel0 = FuncType.parallel
_vectorize0 = FuncType.vectorize
_unroll0 = FuncType.unroll

#Func.realize = lambda x, *a: _realize(x,*a) if not (len(a)==1 and isinstance(a[0], ImageTypes)) else _realize(x,to_dynimage(a[0]))
Func.allow_race_conditions = lambda self: _allow_race_conditions0(self)

Func.gpu_blocks = lambda self, *a: _gpu_blocks0(self, *wrap_gpu_args_int(a))
Func.gpu_single_thread = lambda self, *a: _gpu_single_thread0(self, *wrap_gpu_args_int(a))
Func.gpu_threads = lambda self, *a: _gpu_threads0(self, *wrap_gpu_args_int(a))
Func.gpu_tile = lambda self, *a: _gpu_tile0(self, *wrap_gpu_args_int(a))

Func.serial = lambda self, *a: _serial0(self, *wrap_gpu_args(a))
Func.fuse = lambda self, *a: _fuse0(self, *wrap_gpu_args(a))
Func.split = lambda self, *a: _split0(self, *wrap_gpu_args(a))
Func.tile = lambda self, *a: _tile0(self, *wrap_gpu_args(a))
Func.reorder = lambda self, *a: _reorder0(self, *wrap_gpu_args(a))
Func.bound = lambda self, a, b, c: _bound0(self, a, wrap(b), wrap(c))
Func.parallel = lambda self, *a: _parallel0(self, *wrap_gpu_args(a))
Func.vectorize = lambda self, *a: _vectorize0(self, *wrap_gpu_args_int(a))
Func.unroll = lambda self, *a: _unroll0(self, *wrap_gpu_args_int(a))

#Deprecated
Func.cuda_blocks = lambda self, *a: _gpu_blocks0(self, *wrap_gpu_args_int(a))
Func.cuda_single_thread = lambda self, *a: _gpu_single_thread0(self, *wrap_gpu_args_int(a))
Func.cuda_threads = lambda self, *a: _gpu_threads0(self, *wrap_gpu_args_int(a))
Func.cuda_tile = lambda self, *a: _gpu_tile0(self, *wrap_gpu_args_int(a))

class Func(object):
    """
    A halide function. This class represents one stage in a Halide
    pipeline, and is the unit by which we schedule things. By default
    they are aggressively inlined, so you are encouraged to make lots
    of little functions, rather than storing things in Exprs.

    Constructors::

      Func()      -- Declare a new undefined function with an automatically-generated unique name
      Func(expr)  -- Declare a new function with an automatically-generated unique
                     name, and define it to return the given expression (which may
                     not contain free variables).
      Func(name)  -- Declare a new undefined function with the given name

    """
    def __new__(cls, *args):
        return FuncType(*args)

    def set(self, y):
        """
        Typically one uses f[x, y] = expr to assign to a function. However f.set(expr) can be used also.
        """

    def allow_race_conditions(self):
        """
        Specify that race conditions are permitted for this Func,
        which enables parallelizing over RVars even when Halide cannot
        prove that it is safe to do so. Use this with great caution,
        and only if you can prove to yourself that this is safe, as it
        may result in a non-deterministic routine that returns
        different values at different times or on different machines. */
        """

    def realize(self, x_size=0, y_size=0, z_size=0, w_size=0):
        """
        Evaluate this function over some rectangular domain and return
        the resulting buffer. The buffer should probably be instantly
        wrapped in an Image class.

        One can use f.realize(Buffer) to realize into an existing buffer.
        """

    def compile_to_bitcode(self, filename, list_of_Argument, fn_name=""):
        """
        Statically compile this function to llvm bitcode, with the
        given filename (which should probably end in .bc), type
        signature, and C function name (which defaults to the same name
        as this halide function.
        """

    def compile_to_c(self, filename, list_of_Argument, fn_name=""):
        """
        Statically compile this function to C source code. This is
        useful for providing fallback code paths that will compile on
        many platforms. Vectorization will fail, and parallelization
        will produce serial code.
        """

    def compile_to_file(self, filename_prefix, list_of_Argument, target):
        """
        Various signatures::

          compile_to_file(filename_prefix, list_of_Argument)
          compile_to_file(filename_prefix)
          compile_to_file(filename_prefix, Argument a)
          compile_to_file(filename_prefix, Argument a, Argument b)

        Compile to object file and header pair, with the given
        arguments. Also names the C function to match the first
        argument.
        """

    def compile_jit(self):
        """
        Eagerly jit compile the function to machine code. This
        normally happens on the first call to realize. If you're
        running your halide pipeline inside time-sensitive code and
        wish to avoid including the time taken to compile a pipeline,
        then you can call this ahead of time. Returns the raw function
        pointer to the compiled pipeline.
        """

    def debug_to_file(self, filename):
        """
        When this function is compiled, include code that dumps its values
        to a file after it is realized, for the purpose of debugging.
        The file covers the realized extent at the point in the schedule that
        debug_to_file appears.

        If filename ends in ".tif" or ".tiff" (case insensitive) the file
        is in TIFF format and can be read by standard tools.
        """

    def name(self):
        """
        The name of this function, either given during construction, or automatically generated.
        """

    def value(self):
        """
        The right-hand-side value of the pure definition of this
        function. May be undefined if the function has no pure
        definition yet.
        """

    def dimensions(self):
        """
        The dimensionality (number of arguments) of this
        function. Zero if the function is not yet defined.
        """

    def __getitem__(self, *args):
        """
        Either calls to the function, or the left-hand-side of a
        reduction definition (see \ref RDom). If the function has
        already been defined, and fewer arguments are given than the
        function has dimensions, then enough implicit vars are added to
        the end of the argument list to make up the difference.
        """

    def split(self, old, outer, inner, factor):
        """
        Split a dimension into inner and outer subdimensions with the
        given names, where the inner dimension iterates from 0 to
        factor-1. The inner and outer subdimensions can then be dealt
        with using the other scheduling calls. It's ok to reuse the old
        variable name as either the inner or outer variable.

        The arguments are all Var instances.
        """

    def fuse(self, inner, outer, fused):
        """
        Join two dimensions into a single fused dimension. The fused
         dimension covers the product of the extents of the inner and
         outer dimensions given.
        """

    def parallel(self, var):
        """
        Mark a dimension (Var instance) to be traversed in parallel.
        """

    def vectorize(self, var, factor):
        """
        Split a dimension (Var instance) by the given int factor, then vectorize the
        inner dimension. This is how you vectorize a loop of unknown
        size. The variable to be vectorized should be the innermost
        one. After this call, var refers to the outer dimension of the
        split.
        """

    def unroll(self, var, factor=None):
        """
        Split a dimension (Var instance) by the given int factor, then unroll the inner
        dimension. This is how you unroll a loop of unknown size by
        some constant factor. After this call, var refers to the outer
        dimension of the split.
        """

    def bound(self, min_expr, extent_expr):
        """
        Statically declare that the range over which a function should
        be evaluated is given by the second and third arguments. This
        can let Halide perform some optimizations. E.g. if you know
        there are going to be 4 color channels, you can completely
        vectorize the color channel dimension without the overhead of
        splitting it up. If bounds inference decides that it requires
        more of this function than the bounds you have stated, a
        runtime error will occur when you try to run your pipeline.
        """

    def tile(self, x, y, xo, yo, xi, yi, xfactor, yfactor):
        """
        Traverse in tiled order. Two signatures::

          tile(x, y, xi, yi, xfactor, yfactor)
          tile(x, y, xo, yo, xi, yi, xfactor, yfactor)

        Split two dimensions at once by the given factors, and then
        reorder the resulting dimensions to be xi, yi, xo, yo from
        innermost outwards. This gives a tiled traversal.

        The shorter form of tile reuses the old variable names as
        the new outer dimensions.
        """

    def reorder(self, *args):
        """
        Reorder the dimensions (Var arguments) to have the given nesting
        order, from innermost loop order to outermost.
        """

    def rename(self, old_name, new_name):
        """
        Rename a dimension. Equivalent to split with a inner size of one.
        """

    def gpu_threads(self, *args):
        """
        Tell Halide that the following dimensions correspond to cuda
        thread indices. This is useful if you compute a producer
        function within the block indices of a consumer function, and
        want to control how that function's dimensions map to cuda
        threads. If the selected target is not ptx, this just marks
        those dimensions as parallel.
        """

    def gpu_single_thread(self, *args):
        """
        Tell Halide to run this stage using a single gpu thread and
        block. This is not an efficient use of your GPU, but it can be
        useful to avoid copy-back for intermediate update stages that
        touch a very small part of your Func.
        """

    def gpu_blocks(self, *args):
        """
        Tell Halide that the following dimensions correspond to cuda
        block indices. This is useful for scheduling stages that will
        run serially within each cuda block. If the selected target is
        not ptx, this just marks those dimensions as parallel.
        """

    def gpu(self, block_x, thread_x):
        """
        Three signatures::

          gpu(block_x, thread_x)
          gpu(block_x, block_y, thread_x, thread_y)
          gpu(block_x, block_y, block_z, thread_x, thread_y, thread_z)

        Tell Halide that the following dimensions correspond to cuda
        block indices and thread indices. If the selected target is not
        ptx, these just mark the given dimensions as parallel. The
        dimensions are consumed by this call, so do all other
        unrolling, reordering, etc first.
        """

    def gpu_tile(self, x, x_size):
        """
        Three signatures:

          gpu_tile(x, x_size)
          gpu_tile(x, y, x_size, y_size)
          gpu_tile(x, y, z, x_size, y_size, z_size)

        Short-hand for tiling a domain and mapping the tile indices
        to cuda block indices and the coordinates within each tile to
        cuda thread indices. Consumes the variables given, so do all
        other scheduling first.
        """

    def cuda_threads(self, *args):
        """
        deprecated Old name for #gpu_threads.
        Tell Halide that the following dimensions correspond to cuda
        thread indices. This is useful if you compute a producer
        function within the block indices of a consumer function, and
        want to control how that function's dimensions map to cuda
        threads. If the selected target is not ptx, this just marks
        those dimensions as parallel.
        """

    def cuda_blocks(self, *args):
        """
        deprecated Old name for #cuda_blocks.
        Tell Halide that the following dimensions correspond to cuda
        block indices. This is useful for scheduling stages that will
        run serially within each cuda block. If the selected target is
        not ptx, this just marks those dimensions as parallel.
        """

    def cuda(self, block_x, thread_x):
        """
        deprecated Old name for #cuda.
        Three signatures::

          cuda(block_x, thread_x)
          cuda(block_x, block_y, thread_x, thread_y)
          cuda(block_x, block_y, block_z, thread_x, thread_y, thread_z)

        Tell Halide that the following dimensions correspond to cuda
        block indices and thread indices. If the selected target is not
        ptx, these just mark the given dimensions as parallel. The
        dimensions are consumed by this call, so do all other
        unrolling, reordering, etc first.
        """

    def cuda_tile(self, x, x_size):
        """
        deprecated Old name for #cuda_tile.
        Three signatures:

          cuda_tile(x, x_size)
          cuda_tile(x, y, x_size, y_size)
          cuda_tile(x, y, z, x_size, y_size, z_size)

        Short-hand for tiling a domain and mapping the tile indices
        to cuda block indices and the coordinates within each tile to
        cuda thread indices. Consumes the variables given, so do all
        other scheduling first.
        """

    def reorder_storage(self, *args):
        """
        Scheduling calls that control how the storage for the function
        is laid out. Right now you can only reorder the dimensions.
        """

    def compute_at(self, f, var):
        """
        Compute this function as needed for each unique value of the
        given var (can be a Var or an RVar) for the given calling function f.
        """

    def compute_root(self):
        """
        Compute all of this function once ahead of time.
        """

    def store_at(self, f, var):
        """
        Allocate storage for this function within f's loop over
        var (can be a Var or an RVar). Scheduling storage is optional, and can be used to
        separate the loop level at which storage occurs from the loop
        level at which computation occurs to trade off between locality
        and redundant work.
        """

    def store_root(self):
        """
        Equivalent to Func.store_at, but schedules storage outside the outermost loop.
        """

    def compute_inline(self):
        """
        Aggressively inline all uses of this function. This is the
        default schedule, so you're unlikely to need to call this. For
        a reduction, that means it gets computed as close to the
        innermost loop as possible.
        """

    def update(self):
        """
        Get a handle on the update step of a reduction for the
        purposes of scheduling it. Only the pure dimensions of the
        update step can be meaningfully manipulated (see RDom).
        """

    def function(self):
        """
        Get a handle on the internal halide function that this Func
        represents. Useful if you want to do introspection on Halide
        functions.
        """

# ----------------------------------------------------
# Stage
# ----------------------------------------------------

_allow_race_conditions1 = StageType.allow_race_conditions
_gpu1 = StageType.gpu
_gpu_blocks1 = StageType.gpu_blocks
_gpu_single_thread1 = StageType.gpu_single_thread
_gpu_threads1 = StageType.gpu_threads
_gpu_tile1 = StageType.gpu_tile
_serial1 = StageType.serial
_split1 = StageType.split
_fuse1 = StageType.fuse
_tile1 = StageType.tile
_reorder1 = StageType.reorder
_parallel1 = StageType.parallel
_vectorize1 = StageType.vectorize
_unroll1 = StageType.unroll

#Stage.realize = lambda x, *a: _realize(x,*a) if not (len(a)==1 and isinstance(a[1], ImageTypes)) else _realize(x,to_dynimage(a[1]))
Stage.allow_race_conditions = lambda self: _allow_race_conditions1(self)

Stage.gpu_blocks = lambda self, *a: _gpu_blocks1(self, *wrap_gpu_args_int(a))
Stage.gpu_single_thread = lambda self, *a: _gpu_single_thread1(self, *wrap_gpu_args_int(a))
Stage.gpu_threads = lambda self, *a: _gpu_threads1(self, *wrap_gpu_args_int(a))
Stage.gpu_tile = lambda self, *a: _gpu_tile1(self, *wrap_gpu_args(a))

Stage.serial = lambda self, *a: _serial1(self, *wrap_gpu_args(a))
Stage.fuse = lambda self, *a: _fuse1(self, *wrap_gpu_args(a))
Stage.split = lambda self, *a: _split1(self, *wrap_gpu_args(a))
Stage.tile = lambda self, *a: _tile1(self, *wrap_gpu_args(a))
Stage.reorder = lambda self, *a: _reorder1(self, *wrap_gpu_args(a))
Stage.parallel = lambda self, *a: _parallel1(self, *wrap_gpu_args(a))
Stage.vectorize = lambda self, *a: _vectorize1(self, *wrap_gpu_args_int(a))
Stage.unroll = lambda self, *a: _unroll1(self, *wrap_gpu_args_int(a))

# Deprecated
Stage.cuda_blocks = lambda self, *a: _gpu_blocks1(self, *wrap_gpu_args_int(a))
Stage.cuda_single_thread = lambda self, *a: _gpu_single_thread1(self, *wrap_gpu_args_int(a))
Stage.cuda_threads = lambda self, *a: _gpu_threads1(self, *wrap_gpu_args_int(a))
Stage.cuda_tile = lambda self, *a: _gpu_tile1(self, *wrap_gpu_args(a))



class Stage(object):
    """
    Scheduling calls that control how the domain of this update is
    traversed. This class represents one stage in a Halide
    pipeline, and is the unit by which we schedule things. By default
    they are aggressively inlined, so you are encouraged to make lots
    of little functions, rather than storing things in Exprs.

    """

    def allow_race_conditions(self):
        """
        Specify that race conditions are permitted for this Func,
        which enables parallelizing over RVars even when Halide cannot
        prove that it is safe to do so. Use this with great caution,
        and only if you can prove to yourself that this is safe, as it
        may result in a non-deterministic routine that returns
        different values at different times or on different machines. */
        """

    def split(self, old, outer, inner, factor):
        """
        Split a dimension into inner and outer subdimensions with the
        given names, where the inner dimension iterates from 0 to
        factor-1. The inner and outer subdimensions can then be dealt
        with using the other scheduling calls. It's ok to reuse the old
        variable name as either the inner or outer variable.

        The arguments are all Var instances.
        """

    def fuse(self, inner, outer, fused):
        """
        Join two dimensions into a single fused dimension. The fused
         dimension covers the product of the extents of the inner and
         outer dimensions given.
        """

    def parallel(self, var):
        """
        Mark a dimension (Var instance) to be traversed in parallel.
        """

    def vectorize(self, var, factor):
        """
        Split a dimension (Var instance) by the given int factor, then vectorize the
        inner dimension. This is how you vectorize a loop of unknown
        size. The variable to be vectorized should be the innermost
        one. After this call, var refers to the outer dimension of the
        split.
        """

    def unroll(self, var, factor=None):
        """
        Split a dimension (Var instance) by the given int factor, then unroll the inner
        dimension. This is how you unroll a loop of unknown size by
        some constant factor. After this call, var refers to the outer
        dimension of the split.
        """

    def tile(self, x, y, xo, yo, xi, yi, xfactor, yfactor):
        """
        Traverse in tiled order. Two signatures::

          tile(x, y, xi, yi, xfactor, yfactor)
          tile(x, y, xo, yo, xi, yi, xfactor, yfactor)

        Split two dimensions at once by the given factors, and then
        reorder the resulting dimensions to be xi, yi, xo, yo from
        innermost outwards. This gives a tiled traversal.

        The shorter form of tile reuses the old variable names as
        the new outer dimensions.
        """

    def reorder(self, *args):
        """
        Reorder the dimensions (Var arguments) to have the given nesting
        order, from innermost loop order to outermost.
        """

    def rename(self, old_name, new_name):
        """
        Rename a dimension. Equivalent to split with a inner size of one.
        """

    def gpu_threads(self, *args):
        """
        Tell Halide that the following dimensions correspond to cuda
        thread indices. This is useful if you compute a producer
        function within the block indices of a consumer function, and
        want to control how that function's dimensions map to cuda
        threads. If the selected target is not ptx, this just marks
        those dimensions as parallel.
        """

    def gpu_single_thread(self, *args):
        """
        Tell Halide to run this stage using a single gpu thread and
        block. This is not an efficient use of your GPU, but it can be
        useful to avoid copy-back for intermediate update stages that
        touch a very small part of your Func.
        """

    def gpu_blocks(self, *args):
        """
        Tell Halide that the following dimensions correspond to cuda
        block indices. This is useful for scheduling stages that will
        run serially within each cuda block. If the selected target is
        not ptx, this just marks those dimensions as parallel.
        """

    def gpu(self, block_x, thread_x):
        """
        Three signatures::

          gpu(block_x, thread_x)
          gpu(block_x, block_y, thread_x, thread_y)
          gpu(block_x, block_y, block_z, thread_x, thread_y, thread_z)

        Tell Halide that the following dimensions correspond to cuda
        block indices and thread indices. If the selected target is not
        ptx, these just mark the given dimensions as parallel. The
        dimensions are consumed by this call, so do all other
        unrolling, reordering, etc first.
        """

    def gpu_tile(self, x, x_size):
        """
        Three signatures:

          gpu_tile(x, x_size)
          gpu_tile(x, y, x_size, y_size)
          gpu_tile(x, y, z, x_size, y_size, z_size)

        Short-hand for tiling a domain and mapping the tile indices
        to cuda block indices and the coordinates within each tile to
        cuda thread indices. Consumes the variables given, so do all
        other scheduling first.
        """

    def cuda_threads(self, *args):
        """
        deprecated Old name for #gpu_threads.
        Tell Halide that the following dimensions correspond to cuda
        thread indices. This is useful if you compute a producer
        function within the block indices of a consumer function, and
        want to control how that function's dimensions map to cuda
        threads. If the selected target is not ptx, this just marks
        those dimensions as parallel.
        """

    def cuda_blocks(self, *args):
        """
        deprecated Old name for #cuda_blocks.
        Tell Halide that the following dimensions correspond to cuda
        block indices. This is useful for scheduling stages that will
        run serially within each cuda block. If the selected target is
        not ptx, this just marks those dimensions as parallel.
        """

    def cuda(self, block_x, thread_x):
        """
        deprecated Old name for #cuda.
        Three signatures::

          cuda(block_x, thread_x)
          cuda(block_x, block_y, thread_x, thread_y)
          cuda(block_x, block_y, block_z, thread_x, thread_y, thread_z)

        Tell Halide that the following dimensions correspond to cuda
        block indices and thread indices. If the selected target is not
        ptx, these just mark the given dimensions as parallel. The
        dimensions are consumed by this call, so do all other
        unrolling, reordering, etc first.
        """

    def cuda_tile(self, x, x_size):
        """
        deprecated Old name for #cuda_tile.
        Three signatures:

          cuda_tile(x, x_size)
          cuda_tile(x, y, x_size, y_size)
          cuda_tile(x, y, z, x_size, y_size, z_size)

        Short-hand for tiling a domain and mapping the tile indices
        to cuda block indices and the coordinates within each tile to
        cuda thread indices. Consumes the variables given, so do all
        other scheduling first.
        """


# ----------------------------------------------------
# RDom and RVar
# ----------------------------------------------------

class Argument(object):
    """
    A multi-dimensional domain over which to iterate. Used when
    defining functions as reductions. See apps/bilateral_grid.py for an
    example of a reduction.

    Constructors::

      Argument(String name, Bool is_buffer, Type type)                             -- 1D reduction
    """
    def __new__(cls, *args):
        return ArgumentType(*args)

    def name(self):
        """
        The name of the argument.
        """

    def is_buffer(self, other):
        """
        An argument is either a primitive type (for parameters), or a
        buffer pointer. If 'is_buffer' is true, then 'type' should be
        ignored.
        """

    def type(self):
        """
        If this is a scalar parameter, then this is its type.
        """


class RDom(object):
    """
    A multi-dimensional domain over which to iterate. Used when
    defining functions as reductions. See apps/bilateral_grid.py for an
    example of a reduction.

    Constructors::

      RDom(Expr min, Expr extent, name="")                             -- 1D reduction
      RDom(Expr min0, Expr extent0, Expr min1, Expr extent1, name="")  -- 2D reduction
      (Similar for 3D and 4D reductions)
      RDom(Buffer|ImageParam)                    -- Iterate over all points in the domain

    The following global functions can be used for inline reductions::

        minimum, maximum, product, sum
    """
    def __new__(cls, *args):
        args = [wrap(x) if not isinstance(x,str) else x for x in args]
        return RDomType(*args)

    def defined(self):
        """
        Check if reduction domain is non-NULL.
        """

    def same_as(self, other):
        """
        Check if two reduction domains are the same.
        """

    def dimensions(self):
        """
        Number of dimensions.
        """

    x = property(doc="Access dimension 1 of reduction domain.")
    y = property(doc="Access dimension 2 of reduction domain.")
    z = property(doc="Access dimension 3 of reduction domain.")
    w = property(doc="Access dimension 4 of reduction domain.")

class RVar(object):
    """
    A reduction variable represents a single dimension of a reduction
    domain (RDom). Don't construct them directly, instead construct an
    RDom, and use fields .x, .y, .z, .w to get at the variables. For
    single-dimensional reduction domains, you can just cast a
    single-dimensional RDom to an RVar.
    """
    def __new__(cls, *args):
        return RVarType(*args)

    def min(self):
        """
        The minimum value that this variable will take on.
        """

    def extent(self):
        """
        The number that this variable will take on. The maximum value
        of this variable will be min() + extent() - 1.
        """

    def name(self):
        """
        The name of this reduction variable.
        """

class VarOrRVar(object):
    """
    A class that can represent Vars or RVars. Used for reorder calls
    which can accept a mix of either
    """
    def __new__(cls, *args):
        return VarOrRVarType(*args)


# ----------------------------------------------------
# Image
# ----------------------------------------------------

_flip_xy = True

def flip_xy(flip=True):
    """
    Whether to flip array dimensions when converting to Numpy/PIL.

    Halide examples use f[x,y] convention. If flip_xy(True) is called (the default) then the variables
    are flipped to I[y,x] convention when converting to Numpy/PIL.
    """
    global _flip_xy
    _flip_xy = flip

import ctypes

def _image_getattr(self, name):
    if name == '__array_interface__':
        #print('get array')
        D = self
        t = D.type()
        if t.is_int():
            typestr = '|i%d'%(t.bits/8)
        elif t.is_uint():
            typestr = '|u%d'%(t.bits/8)
        elif t.is_float():
            typestr = '|f%d'%(t.bits/8)
        else:
            raise ValueError('Unknown type %r'%t)
        shape = tuple([D.extent(i) for i in range(D.dimensions())])
        strides = tuple([D.stride(i)*(t.bits/8) for i in range(D.dimensions())])
        if _flip_xy and len(strides) >= 2:
            strides = (strides[1], strides[0]) + strides[2:]
            shape = (shape[1], shape[0]) + shape[2:]

        #print("About to call image_to_string on", type(self))
        data = image_to_string(self)
        #print(type(data))
        #data = image_to_uint8(self)
        #data_p = get_image_bytes(self)
        #print(type(data))
        #data = ctypes.c_uint8_p(int(data_p))
        #void_data_p = ctypes.c_void_p(int(data_p))
        #uint8_data = ctypes.cast(void_data_p, ctypes.POINTER(ctypes.c_uint8))

        return {'shape': shape,
                'typestr': typestr,
                #'data': uint8_data,
                'data': data,
                'strides': strides}
    raise AttributeError(name)

for _ImageT in ImageTypes:
    _ImageT.__getitem__ = _generic_getitem_expr

Image_int8.type = lambda x: Int(8)
Image_int16.type = lambda x: Int(16)
Image_int32.type = lambda x: Int(32)
Image_uint8.type = lambda x: UInt(8)
Image_uint16.type = lambda x: UInt(16)
Image_uint32.type = lambda x: UInt(32)
Image_float32.type = lambda x: Float(32)
Image_float64.type = lambda x: Float(64)

def _to_pil(I, maxval=None):
    "Convert an Image instance to a PIL image."
    #print(I.__array_interface__)
    A = numpy.asarray(I)
    if maxval is not None:
        A = numpy.asarray(A,'float32')*(1.0/maxval)
#     print('converted to numpy', A.dtype)
    if A.dtype == numpy.uint8:
        pass
    elif A.dtype == numpy.uint16: #numpy.issubdtype(A.dtype, 'uint16'):
        A = numpy.array(A/256, 'uint8')
    elif A.dtype == numpy.uint32: #numpy.issubdtype(A.dtype, 'uint32'):
        A = numpy.array(A/256**3, 'uint8')
    elif A.dtype == numpy.float32 or A.dtype == numpy.float64: #numpy.issubdtype(A.dtype, 'float32') or numpy.issubdtype(A.dtype, 'float64'):
        A = numpy.array(A*255.0, 'uint8')
    else:
        raise ValueError('Unsupported dtype %r' % A.dtype)
    if len(A.shape) == 3 and A.shape[2] == 1:
        A = A[:,:,0]
    return PIL.fromarray(A)

def _show_image(I, maxval=None):
    _to_pil(I, maxval).show()

def _save_image(I, filename, maxval=None):
    _to_pil(I, maxval).save(filename)

for _ImageT in ImageTypes:
    _ImageT.save = lambda *args, **kw: _save_image(*args, **kw)
    _ImageT.to_pil = lambda *args, **kw: _to_pil(*args, **kw)
    _ImageT.set = _generic_set
    _ImageT.__getattr__ = _image_getattr
    _ImageT.show = lambda *args, **kw: _show_image(*args, **kw)
    _ImageT.tostring = lambda self: image_to_string(self)

def _numpy_to_image(a, dtype, C):
    a = numpy.asarray(a, dtype)
    shape = a.shape
    strides = a.strides
    if _flip_xy and len(shape) >= 2:
        shape = (shape[1], shape[0]) + shape[2:]
        strides = (strides[1], strides[0]) + strides[2:]

    ans = C(*shape)
    assign_array(ans, a.__array_interface__['data'][0], *strides)
    return ans

def _numpy_to_type(a):
    d = {numpy.dtype('int8'): Int(8),
         numpy.dtype('int16'): Int(16),
         numpy.dtype('int32'): Int(32),
         numpy.dtype('uint8'): UInt(8),
         numpy.dtype('uint16'): UInt(16),
         numpy.dtype('uint32'): UInt(32),
         numpy.dtype('float32'): Float(32),
         numpy.dtype('float64'): Float(64)}
    return d[a.dtype]

class Image(object):
    """
    Construct an Image::

        Image(contents, [scale=None])
        Image([typeval=Int(n), UInt(n), Float(n), Bool()], contents, [scale=None])

    The contents can be:

        - PIL image
        - Numpy array
        - Filename of existing file (typeval defaults to UInt(8))
        - halide.Buffer
        - An int or tuple -- constructs an n-D image (typeval argument is required).

    The image can be indexed via I[x], I[y,x], etc, which gives a Halide Expr.

    If not provided (or None) then the typeval is inferred from the input argument.

    For PIL, numpy, and filename constructors, if scale is provided then the input is scaled by the floating point scale factor
    (for example, Image(filename, UInt(16), 1.0/256) reads a UInt(16) image rescaled to have maximum value 255). If omitted,
    scale is set to convert between source and target data type ranges, where int types range from 0 to maxval, and float types
    range from 0 to 1.
    """
    def __new__(cls, typeval, contents=None, scale=None):
        if contents is None:                        # If contents is None then Image(contents) is assumed as the call
            (typeval, contents) = (None, typeval)

        if isinstance(contents, str_types):    # Convert filename to PIL image
            contents = PIL.open(contents)
            if typeval is None:
                typeval = UInt(8)
                contents = numpy.asarray(contents, 'uint8')

        if hasattr(contents, 'putpixel'):           # Convert PIL images to numpy
            contents = numpy.asarray(contents)

        if typeval is None:
            if hasattr(contents, 'type'):
                typeval = contents.type()
            elif isinstance(contents, numpy.ndarray):
                typeval = _numpy_to_type(contents)
            elif isinstance(contents, Realization):
                typeval = contents.as_vector()[0].type()
            else:
                raise ValueError('unknown halide.Image constructor %r' % contents)

        assert isinstance(typeval, TypeType), typeval
        sig = (typeval.bits, typeval.is_int(), typeval.is_uint(), typeval.is_float())

        if sig == (8, True, False, False):
            C = Image_int8
            target_dtype = 'int8'
        elif sig == (16, True, False, False):
            C = Image_int16
            target_dtype = 'int16'
        elif sig == (32, True, False, False):
            C = Image_int32
            target_dtype = 'int32'
        elif sig == (8, False, True, False):
            C = Image_uint8
            target_dtype = 'uint8'
        elif sig == (16, False, True, False):
            C = Image_uint16
            target_dtype = 'uint16'
        elif sig == (32, False, True, False):
            C = Image_uint32
            target_dtype = 'uint32'
        elif sig == (32, False, False, True):
            C = Image_float32
            target_dtype = 'float32'
        elif sig == (64, False, False, True):
            C = Image_float64
            target_dtype = 'float64'
        else:
            raise ValueError('unimplemented halide.Image type signature %r' % typeval)

        if isinstance(contents, numpy.ndarray):
            if scale is None:
                in_range = _numpy_to_type(contents).typical_max()
                out_range = typeval.typical_max()
                if in_range != out_range:
                    scale = float(out_range)/in_range
            if scale is not None:
                contents = numpy.asarray(numpy.asarray(contents,'float')*float(scale), target_dtype)
            return _numpy_to_image(contents, target_dtype, C)
        elif isinstance(contents, ImageTypes+(ImageParamType,BufferType,Realization)):
            return C(contents)
        elif isinstance(contents, tuple) or isinstance(contents, list) or isinstance(contents, integer_types):
            if isinstance(contents, integer_types):
                contents = (contents,)
            if not all(isinstance(x, integer_types) for x in contents):
                raise ValueError('halide.Image constructor did not receive a tuple of ints for Image size')
            return C(*contents)
        else:
            raise ValueError('unknown Image constructor contents %r' % contents)

    def show(self, maxval=None):
        """
        Shows an Image instance on the screen, by converting through numpy and PIL.

        If maxval is not None then rescales so that bright white is equal to maxval.
        """

    def save(self, filename, maxval=None):
        """
        Save an Image (converted via PIL).

        If maxval is not None then rescales so that bright white is equal to maxval.
        """

    def to_pil(self):
        """
        Convert to PIL (Python Imaging Library) image.

        The fromarray() constructor in PIL does not work due to ignoring array stride so this is a workaround.
        """

    def tostring(self):
        """
        Convert to str.
        """

    def set(self, contents):
        """
        Sets contents of Image or ImageParam::

            set(Image, Buffer)
            set(ImageParam, Buffer|Image)
        """

    def channels(self):
        """
        Number of channels.
        """

    def copy_to_host(self):
        """
        Manually copy-back data to the host, if it's on a device. This
        is done for you if you construct an image from a buffer, but
        you might need to call this if you realize a gpu kernel into an
        existing image.
        """

    def defined(self):
        """
        Return whether buffer points to actual data (non-NULL Image).
        """

    def dimensions(self):
        """
        Get the dimensionality of the data. Typically two for grayscale images, and three for color images.
        """

    def extent(self, i):
        """
        Extent of dimension i.
        """

    def height(self):
        """
        Get the extent of dimension 0, which by convention we use as
        the height of the image.
        """

    def set_host_dirty(self, dirty=True):
        """
        Mark the buffer as dirty-on-host.  is done for you if you
        construct an image from a buffer, but you might need to call
        this if you realize a gpu kernel into an existing image, or
        modify the data via some other back-door.
        """

    def stride(self, dim):
        """
        Get the number of elements in the buffer between two adjacent
        elements in the given dimension. For example, the stride in
        dimension 0 is usually 1, and the stride in dimension 1 is
        usually the extent of dimension 0. This is not necessarily true
        though.
        """

    def width(self):
        """
        Get the extent of dimension 0, which by convention we use as
        the width of the image.
        """

    def type(self):
        """
        Return Type instance for the data type of the image.
        """

# ----------------------------------------------------
# Param
# ----------------------------------------------------

class Param(object):
    """
    A scalar parameter to a halide pipeline. If you're jitting, this
    should be bound to an actual value of type T using the set method
    before you realize the function uses this. If you're statically
    compiling, this param should appear in the argument list.

    Constructors::

      Param(typeval, [value])
      Param(typeval, name, [value])

    See Type for how to construct typeval.
    """
    def __new__(cls, typeval, *args):
        assert isinstance(typeval, TypeType)
        sig = (typeval.bits, typeval.is_int(), typeval.is_uint(), typeval.is_float())
        if sig == (8, True, False, False):
            C = Param_int8
        elif sig == (16, True, False, False):
            C = Param_int16
        elif sig == (32, True, False, False):
            C = Param_int32
        elif sig == (8, False, True, False):
            C = Param_uint8
        elif sig == (16, False, True, False):
            C = Param_uint16
        elif sig == (32, False, True, False):
            C = Param_uint32
        elif sig == (32, False, False, True):
            C = Param_float32
        elif sig == (64, False, False, True):
            C = Param_float64
        else:
            raise ValueError('unimplemented Param type signature %r' % typeval)
        if len(args) == 1:          # Handle special cases since SWIG apparently cannot convert int32_t to int.
            if isinstance(args[0], (int, float)):
                ans = C()
                set(ans, args[0])
                return ans
        elif len(args) == 2:
            if isinstance(args[1], (int, float)):
                ans = C(args[0])
                set(ans, args[1])
                return ans
        return C(*args)

    def name(self):
        """
        String name of parameter.
        """

    def get(self):
        """
        Get the current value of this parameter. Only meaningful when jitting.
        """

    def set(self, val):
        """
        Set the current value of this parameter. Only meaningful when jitting.
        """

    def type(self):
        """
        Get the Halide type.
        """

for ParamT in ParamTypes: # + (DynUniform,):
    ParamT.set = lambda x, y: set(x, y) #_generic_assign

# ----------------------------------------------------
# Various image types
# ----------------------------------------------------

#UniformImage.__setitem__ = lambda x, key, value: assign(call(x, *[wrap(y) for y in key]), wrap(value)) if isinstance(key,tuple) else assign(call(x, key), wrap(value))

for _ImageT in [ImageParamType]:
    _ImageT.__getitem__ = _generic_getitem_expr
    _ImageT.set = lambda x, y: set(x, Image(y) if (isinstance(y,numpy.ndarray) or hasattr(y, 'putpixel')) else y)
    #_ImageT.save = lambda x, y: save_png(x, y)

# ----------------------------------------------------
# Type
# ----------------------------------------------------

class Type(object):
    """
    A Halide type. Constructors are the following global functions::

      Float(nbits)
      UInt(nbits)
      Int(nbits)
      Bool([width])
    """

    def __new__(cls, *args):
        return TypeType(*args)

    def to_numpy(self):
        """
        Convert to a numpy dtype instance.
        """

    def typical_max(self):
        """
        Get typical maximum value of the type (1.0 for float, otherwise max int value).
        """

    def is_int(self):
        """
        True if Int() type.
        """

    def is_uint(self):
        """
        True if UInt() type.
        """

    def is_float(self):
        """
        True if Float() type.
        """

    def is_bool(self):
        """
        True if Bool() type.
        """

    def imin(self):
        """
        Minimum value for integer (assert if not integer).
        """

    def imax(self):
        """
        Maximum value for integer (assert if not integer).
        """

    bits = property(doc='Bits of type')

def _type_typical_max(typeval):          # The typical maximum value used for image processing (1.0 for float types)
    if typeval.is_uint():
        return 2**(typeval.bits)-1
    elif typeval.is_int():
        return 2**(typeval.bits-1)-1
    elif typeval.is_float():
        return 1.0
    else:
        raise ValueError('unknown typeval %r'%typeval)

def _type_to_numpy(typeval):
    if typeval.is_int():
        return numpy.dtype('int%d'%typeval.bits)
    elif typeval.is_uint():
        return numpy.dtype('uint%d'%typeval.bits)
    elif typeval.is_float():
        return numpy.dtype('float%d'%typeval.bits)
    else:
        raise ValueError('unknown type %r'%typeval)

TypeType.typical_max = _type_typical_max
TypeType.to_numpy = _type_to_numpy

# ----------------------------------------------------
# Repr
# ----------------------------------------------------

VarType.__repr__ = lambda self: 'Var(%r)'%self.name()
TypeType.__repr__ = lambda self: 'UInt(%d)'%self.bits if self.is_uint() else 'Int(%d)'%self.bits if self.is_int() else 'Float(%d)'%self.bits if self.is_float() else 'Type(%d)'%self.bits
ExprType.__repr__ = lambda self: 'Expr(%s)' % ', '.join([repr(self.type())]) #  + [str(_x) for _x in self.vars()]
FuncType.__repr__ = lambda self: 'Func(%r)' % self.name()

# ----------------------------------------------------
# Global functions
# ----------------------------------------------------

# Wrap unary math functions

for _funcname in 'sqrt sin cos exp log floor ceil round asin acos atan tan sinh asinh cosh acosh tanh atanh fast_log fast_exp abs'.split():
    exec("_%s = %s" % (_funcname, _funcname))
    exec("%s = lambda x: _%s(wrap(x))" % (_funcname, _funcname))

# Wrap binary math functions
for _funcname in 'max min hypot fast_pow'.split():
    exec("_%s = %s" % (_funcname, _funcname))
    exec("%s = lambda x, y: _%s(wrap(x), wrap(y))" % (_funcname, _funcname))

#_debug = debug         # Debug node was removed from C++ layer
#debug  = lambda x, y, *a: _debug(wrap(x), y, *a)

_clamp = clamp
_cast = cast
_select = select

pow    = lambda x, y: wrap(x)**wrap(y)

clamp  = lambda x, y, z: _clamp(wrap(x), wrap(y), wrap(z))
cast   = lambda x, y: _cast(x, wrap(y))
select = lambda x, y, z: _select(wrap(x), wrap(y), wrap(z))

minimum = lambda x: minimum_func(wrap(x))
maximum = lambda x: maximum_func(wrap(x))
product = lambda x: product_func(wrap(x))
sum     = lambda x: sum_func(wrap(x))

lambda_counter=[0]

def lambda0D(exp):
    """
    Create a 0 dimensional halide function that returns the given
    expression. The function may have more dimensions if the expression
    contains implicit arguments.
    Replacement for halide_lambda
    """
    prev = Func('prev%d'%lambda_counter[0])
    lambda_counter[0] += 1
    prev = exp
    return prev

def lambda2D(x, y, exp):
    """
    Create a 2 dimensional halide function that returns the given
    expression. The function may have more dimensions if the expression
    contains implicit arguments.
    Replacement for halide_lambda
    """
    prev = Func('prev%d'%lambda_counter[0])
    lambda_counter[0] += 1
    prev[x, y] = exp
    return prev

def lambda3D(x, y, z, exp):
    """
    Create a 3 dimensional halide function that returns the given
    expression. The function may have more dimensions if the expression
    contains implicit arguments.
    Replacement for halide_lambda
    """
    prev = Func('prev%d'%lambda_counter[0])
    lambda_counter[0] += 1
    prev[x, y, z] = exp
    return prev

def lambda4D(x, y, z, w, exp):
    """
    Create a 4 dimensional halide function that returns the given
    expression. The function may have more dimensions if the expression
    contains implicit arguments.
    Replacement for halide_lambda
    """
    prev = Func('prev%d'%lambda_counter[0])
    lambda_counter[0] += 1
    prev[x, y, z, w] = exp
    return prev

def lambda5D(x, y, z, w, v, exp):
    """
    Create a 5 dimensional halide function that returns the given
    expression. The function may have more dimensions if the expression
    contains implicit arguments.
    Replacement for halide_lambda
    """
    prev = Func('prev%d'%lambda_counter[0])
    lambda_counter[0] += 1
    prev[x, y, z, w, v] = exp
    return prev


# ----------------------------------------------------
# Constructors
# ----------------------------------------------------

class Expr(object):
    """
    An expression or fragment of Halide code.

    One can explicitly coerce most types to Expr via the Expr(x) constructor.
    The following operators are implemented over Expr, and also other types
    such as Image, Func, Var, RVar generally coerce to Expr when used in arithmetic::

      + - * / % ** & |
      -(unary) ~(unary)
      < <= == != > >=
      += -= *= /=

    The following math global functions are also available::

       Unary:
         abs acos acosh asin asinh atan atanh ceil cos cosh exp
         fast_exp fast_log floor log round sin sinh sqrt tan tanh

       Binary:
         hypot fast_pow max min pow

       Ternary:
         clamp(x, lo, hi)                  -- Clamp expression to [lo, hi]
         select(cond, if_true, if_false)   -- Return if_true if cond else if_false
    """
    def __new__(cls, *args):
        #return wrap(*args)
        w = wrap(*args)
        return w

    def type(self):
        """
        Type of expression.
        """

class ImageParam(object):
    """
    An Image parameter to a halide pipeline. E.g., the input image.

    Constructor::

      ImageParam(Type t, int dims, name="")

    The image can be indexed via I[x], I[y,x], etc, which gives a Halide Expr. Supports most of
    the methods of Image.
    """
    def __new__(cls, *args):
        return ImageParamType(*args)

    def name(self):
        """
        Get name of ImageParam.
        """

    def set(self, I):
        """
        Bind a Buffer, Image, numpy array, or PIL image. Only relevant for jitting.
        """

    def get(self):
        """
        Get the Buffer that is bound. Only relevant for jitting.
        """

# ----------------------------------------------------

# Work around swig bug for python3
# see https://github.com/swig/swig/pull/136
# also used in python2 since using __future__.division
#for C in [VarType, FuncRefExpr, FuncRefVar]:
for C in (ExprType, FuncRefExpr, FuncRefVar, VarType, RDomType, RVarType, FuncType) + ParamTypes:
    C.__truediv__ = C.__div__
    C.__itruediv__ = C.__idiv__

# ----------------------------------------------------
# Test
# ----------------------------------------------------

import time

def test_core():
    input = ImageParam(UInt(16),2)
    var_x = Var('x')
    var_y = Var('y')
    blur_x = Func('blur_x')
    blur_y = Func('blur_y')
    expr_x = Expr(var_x)
    #print(expr_x)
    expr_y = Expr(var_y)
    #Expr(x) + Expr(y)

    T0 = time.time()
    n = 1

    def check(L):
        for x in L:
            assert isinstance(x, ExprType), x

    for i in range(n):
        for x in [expr_x, var_x, 1, 1.3]:
            for y in [expr_y, var_y, 1, 1.3]:
                if isinstance(x, (int,float)) and isinstance(y, (int, float)):
                    continue
                #print type(x), type(y)
                check([x + y,
                x - y,
                x * y,
                x / y,
                x % y,
                x < y,
                x <= y,
                x == y,
                x != y,
                x > y,
                x >= y,
                -x] + [cast(Bool(), x) & cast(Bool(), y),
                cast(Bool(), x) | cast(Bool(), y)] + [~cast(Bool(), x)] if not isinstance(x, (int, float)) else [])

                if isinstance(x, ExprType):
                    x += y
                    check([x])
                    x -= y
                    check([x])
                    x *= y
                    check([x])
                    #x /= y
                    x.__idiv__(y)
                    check([x])

    #blur_x += expr_x
    #blur_x(expr_x)

    T1 = time.time()

    x = var_x
    y = var_y
    z = Var()
    q = Var()

    assert isinstance(x+1, ExprType)
    assert isinstance(x/y, ExprType)
    assert isinstance((x/y)+(x-1), ExprType)
    assert isinstance(blur_x[x-1,y], FuncRefExprType)
    assert isinstance(blur_x[x, y], FuncRefVarType)
    assert isinstance(blur_x[x-1], FuncRefExprType)
    assert isinstance(blur_x[x-1,y,z], FuncRefExprType)
    assert isinstance(blur_x[x-1,y,z,q], FuncRefExprType)
    f = Func()
    f[x,y]=x+1

    print('halide.test_core:                    OK')

def func_varlist(f):
    return [x.name() for x in f.args()]

def get_blur():
    input = ImageParam(UInt(16), 3, 'input')
    x = Var('x')
    y = Var('y')
    c = Var('c')
    input_clamped = Func('input_clamped')
    blur_x = Func('blur_x')
    blur_y = Func('blur_y')

    input_clamped[x,y,c] = input[clamp(Expr(x),cast(Int(32),Expr(0)),cast(Int(32),Expr(input.width()-1))),
                                 clamp(Expr(y),cast(Int(32),Expr(0)),cast(Int(32),Expr(input.height()-1))),
                                 c] #clamp(Expr(c),Expr(0),Expr(input.width()-1))]

    bb_a = input_clamped[x-1,y,c] / 4
    bb_b = input_clamped[x,y,c] / 4
    bb_c = input_clamped[x+1,y,c] / 4
    blur_x[x,y,c] = (bb_a + bb_b + bb_c) / 3
    blur_y[x,y,c] = ((blur_x[x,y-1,c] + blur_x[x,y,c] + blur_x[x,y+1,c]) / 3) * 4

    return (input, x, y, c, blur_x, blur_y, input_clamped)

def builtin_image(filename=None, raw_string=False):
    "PIL Image instance for one of the built-in input images (in halide/data/ or installed with the package)."
    if filename is None:
        filename = 'rgb.png'
    s = pkgutil.get_data('halide', os.path.join('data', filename))
    if raw_string:
        return s

    if sys.version < '3':
        f = cStringIO.StringIO(s)
    else:
        f = io.BytesIO(s)
    return PIL.open(f)

def filter_image(input, out_func, in_image, disp_time=False, compile=True, eval_func=None, out_dims=None, times=1): #, pad_multiple=1):
    """
    Utility function to filter an image filename or numpy array with a Halide Func, returning output Image of the same size.

    Given input and output Funcs, and filename/numpy/PIL array (in_image), returns evaluate. Calling evaluate() returns the output Image.
    """
#    print('filter_image, input=', input, 'dtype', input.type())
    dtype = input.type()
    if isinstance(input, ImageParamType):
        if isinstance(in_image, str):
            input_png = Image(dtype, in_image)
        elif isinstance(in_image, numpy.ndarray):
            input_png = Image(in_image)
        elif hasattr(in_image, 'putpixel'):         # PIL image
            input_png = Image(dtype, in_image)
        else:
            input_png = in_image
        input.set(input_png)
    else:
        input_png = input

    w = input_png.width() if out_dims is None else out_dims[0]
    h = input_png.height() if out_dims is None else out_dims[1]
    nchan = input_png.channels() if out_dims is None else (out_dims[2] if len(out_dims) >= 3 else 1)
    #print(w, h, nchan, out_dims)
    if compile:
        out_func.compile_jit()

    def evaluate():
        T = []
        for i in range(times):
            T0 = time.time()
            if eval_func is not None:
                realized = eval_func(input_png)
            else:
                if nchan != 1:
                    realized = out_func.realize(w, h, nchan)
                else:
                    realized = out_func.realize(w, h)
            T.append(time.time()-T0)
        out = Image(realized) #.set(realized)


        assert out.width() == w and out.height() == h, (out.width(), out.height(), w, h)
        if disp_time:
            if times > 1:
                print('Filtered %d times, min: %.6f secs, mean: %.6f secs.' % (times, numpy.min(T), numpy.mean(T)))
            else:
                print('Filtered in %.6f secs' % T[0])

        return out
    return evaluate

# -------------------------------------------------------------------------------------------
# Unit Tests
# -------------------------------------------------------------------------------------------

def example_out():
    (input, x, y, c, blur_x, blur_y, input_clamped) = get_blur()

    blur_y.unroll(y,16).vectorize(x,16)
    blur_y.compile_jit()

    return filter_image(input, blur_y, builtin_image())()

def test_blur():
    (input, x, y, c, blur_x, blur_y, input_clamped) = get_blur()

    for f in [blur_y, blur_x]:
        assert len(f.args()) == 3
        args = f.args()
        assert args[0].name()=='x'
        assert args[1].name()=='y'
        assert args[2].name()=='c'
    str(blur_y.value())

    #assert len(blur_y.value().funcs()) == 1, list(blur_y.value().funcs())
    #assert set(all_funcs(blur_y).keys()) == set(['blur_x', 'blur_y', 'input_clamped']), set(all_funcs(blur_y).keys())
    assert func_varlist(blur_y) == ['x', 'y', 'c'], func_varlist(blur_y)
    assert func_varlist(blur_x) == ['x', 'y', 'c'], func_varlist(blur_x)

    for i in range(2):
        input_arg = Argument('input', True, UInt(16))
        blur_y.compile_to_file('halide_blur', input_arg)
        try:    os.remove('halide_blur.o')
        except: pass
        try:    os.remove('halide_blur.h')
        except: pass
#        blur_y.compile_to_file('halide_blur', [input_arg])

        outf = filter_image(input, blur_y, builtin_image())
        T0 = time.time()
        out = outf()
        T1 = time.time()
        #print(T1-T0, 'secs')

        out_filename = 'test_out.png'
        out.save(out_filename)
        s = image_to_string(out)
        assert isinstance(s, (str, bytes))
        I1 = numpy.asarray(PIL.open(out_filename))
        os.remove(out_filename)

    print('halide.test_blur:                    OK')

def test_func(compile=True, in_image=None):
    (input, x, y, c, blur_x, blur_y, input_clamped) = get_blur()

    if in_image is None:
        in_image = builtin_image()
    outf = filter_image(input, blur_y, in_image, compile=compile)
    out = [None]

    def test(func):
        T0 = time.time()
        out[0] = outf()
        T1 = time.time()
        return T1-T0

    return locals()

def test_segfault():
    locals_d = test_func(compile=False)
    x = locals_d['x']
    y = locals_d['y']
    c = locals_d['c']
    blur_y = locals_d['blur_y']
    test = locals_d['test']

    exit_on_signal()
    # TODO: Implement a test that segfaults...

def test_examples():
    """
    This function seems deprecated
    """
    in_grayscale = '../../apps/images/gray.png'
    in_color = '../../apps/images/rgb.png'

    names = []
    do_filter = True
    from importlib import import_module

#    for example_name in ['interpolate']: #
    for example_name in 'interpolate blur dilate local_laplacian'.split(): #[examples.snake, examples.blur, examples.dilate, examples.boxblur_sat, examples.boxblur_cumsum, examples.local_laplacian]:
#    for example_name in 'interpolate blur dilate boxblur_sat boxblur_cumsum local_laplacian snake'.split(): #[examples.snake, examples.blur, examples.dilate, examples.boxblur_sat, examples.boxblur_cumsum, examples.local_laplacian]:
#    for example_name in 'interpolate snake blur dilate boxblur_sat boxblur_cumsum local_laplacian'.split(): #[examples.snake, examples.blur, examples.dilate, examples.boxblur_sat, examples.boxblur_cumsum, examples.local_laplacian]:

        example = import_module("apps." + example_name)
        first = True
#    for example in [apps.boxblur_cumsum]:
        for input_image0 in [in_grayscale, in_color]:
            for dtype in [UInt(8), UInt(16), UInt(32), Float(32), Float(64)]:
                input_image = input_image0
#            for dtype in [UInt(16)]:
#            for dtype in [UInt(8), UInt(16)]:
                if example_name == "local_laplacian":
                    if input_image == in_color and (dtype == UInt(8) or dtype == UInt(16)):
                        pass
                    else:
                        continue
                #print(dtype.isFloat(), dtype.bits)
                if example_name == "interpolate":
                    if not dtype.is_float() or input_image != in_color:
                        continue
                    input_image = '../../apps/images/rgba.png'
        #        (in_func, out_func) = apps.blur_color(dtype)
    #            (in_func, out_func) = apps.blur(dtype)
                (in_func, out_func, eval_func, scope) = example(dtype)
                if first:
                    first = False
                    names.append(example_name)
        #        print('got func')
        #        outf = filter_image(in_func, out_func, in_filename, dtype)
        #        print('got filter')
                if do_filter:
                    outf = filter_image(in_func, out_func, input_image, disp_time=True, eval_func=eval_func)
                    out = outf()
                    out.show()
                    A = numpy.asarray(out)
        #        print('shown')
#                print(numpy.min(A.flatten()), numpy.max(A.flatten()))
        #        out.save('out.png')
#    print()
#    print('Function names:')
#    for (example_name, func_names) in names:
#        print example_name, func_names

def test_numpy():
    def dist(a,b):
        a = numpy.asarray(a, 'float64') - numpy.asarray(b, 'float64')
        a = a.flatten()
        return numpy.mean(numpy.abs(a))

    for dtype in [UInt(8), UInt(16), UInt(32), Float(32), Float(64)]:
        for mul in [0, 1]:
            a = numpy.asarray(Image(dtype, builtin_image()))*mul
            b = Image(a)
            c = numpy.asarray(b)
            assert a.dtype == c.dtype
            assert dist(a,c) < 1e-8
            if dtype == UInt(8):
                d = numpy.asarray(b.to_pil())
                e = numpy.asarray(PIL.fromarray(c))
                assert dist(a,d) < 1e-8
                assert dist(a,e) < 1e-8

            if dtype == UInt(16):
                locals_d = test_func(in_image=a)
                test = locals_d['test']
                blur_y = locals_d['blur_y']
                out = locals_d['out']
                test(blur_y)
                c = numpy.asarray(out[0])
                assert dist(a,c)<=1500, dist(a,c)

    print('halide.test_numpy:                   OK')

def test_minimal():
    f1 = Func()
    f2 = Func()
    f3 = Func()
    f4 = Func()
    x = Var()
    y = Var()
    z = Var()
    w = Var()
    f1[x] = 1.0
    f2[x,y] = x
    f3[x,y,z] = y
    f4[x,y,z,w] = 1.0
    print('halide.test_minimal:                 OK')

def test_image_constructors():
    def check(out, typeval_p):
        if typeval_p is None:
            assert isinstance(out, ImageTypes)
        else:
            if typeval_p == Int(8):
                assert isinstance(out, Image_int8)
            elif typeval_p == Int(16):
                assert isinstance(out, Image_int16)
            elif typeval_p == Int(32):
                assert isinstance(out, Image_int32)
            elif typeval_p == UInt(8):
                assert isinstance(out, Image_uint8)
            elif typeval_p == UInt(16):
                assert isinstance(out, Image_uint16)
            elif typeval_p == UInt(32):
                assert isinstance(out, Image_uint32)
            elif typeval_p == Float(32):
                assert isinstance(out, Image_float32)
            elif typeval_p == Float(64):
                assert isinstance(out, Image_float64)
            else:
                raise ValueError

    I = builtin_image('rgb.png')
    A = numpy.asarray(I)

    filename = tempfile.mktemp(suffix='.png')
    try:
        f = open(filename, 'wb')
        f.write(builtin_image('rgb.png', raw_string=True))
        f.close()

        I8 = Image(UInt(8), I)
        I16 = Image(UInt(16), I)
        Ip = ImageParam(UInt(16), 2)

        check(Image(I), UInt(8))
        check(Image(A), UInt(8))
        B = Buffer(UInt(8), 5, 5, 3)

        check(Image(B), UInt(8))
#        check(Image(I8), UInt())
#        check(Image(I16))
#        check(Image(Ip))

        for typeval in [None, Int(8), UInt(16), UInt(32), Float(64)]:
            check(Image(typeval, I), typeval)
            check(Image(typeval, A), typeval)
            check(Image(typeval, filename), typeval)
            if sys.version < '3':
                # python 2
                check(Image(typeval, unicode(filename)), typeval)
            else:
                # python 3
                #check(Image(typeval, bytes(filename, "utf-8")), typeval)
                pass

            if typeval is not None:
                B = Buffer(typeval, 5, 5, 3)
                check(Image(typeval, B), typeval)
            #check(Image(typeval, I8))
            #check(Image(typeval, I16))
            #check(Image(typeval, Ip))
            if typeval is not None:
                check(Image(typeval, 10), typeval)
                check(Image(typeval, (10, 10)), typeval)
                check(Image(typeval, (10, 10, 1)), typeval)
                check(Image(typeval, (10, 10, 3)), typeval)
                check(Image(typeval, (10, 10, 4)), typeval)
    finally:
        try:
            os.remove(filename)
        except:
            pass

    print('halide.test_image_constructors:      OK')

def test():
    exit_on_signal()

    test_minimal()
    test_blur()
    test_core()
    test_numpy()
    test_image_constructors()



if __name__ == '__main__':
    test()

