def get_put_kernel(context, dtype, idx_dtype, vec_count=1):
    ctx = {
            "idx_tp": dtype_to_ctype(idx_dtype),
            "tp": dtype_to_ctype(dtype),
            }

    args = [
            VectorArg(dtype, "dest%d" % i, with_offset=True)
            for i in range(vec_count)
            ] + [
                VectorArg(idx_dtype, "gmem_dest_idx", with_offset=True),
            ] + [
                VectorArg(dtype, "src%d" % i, with_offset=True)
                for i in range(vec_count)
            ] + [
                VectorArg(np.uint8, "use_fill", with_offset=True)
            ] + [
                VectorArg(np.int64, "val_ary_lengths", with_offset=True)
            ]

    body = (
            "%(idx_tp)s dest_idx = gmem_dest_idx[i];\n" % ctx
            + "\n".join(
                    "dest{i}[dest_idx] = (use_fill[{i}] ? src{i}[0] : "
                    "src{i}[i % val_ary_lengths[{i}]]);".format(i=i)
                    for i in range(vec_count)
                    )
            )

    return get_elwise_kernel(context, args, body,
            preamble=dtype_to_c_struct(context.devices[0], dtype),
            name="put")

def get_weighted_inner_kernel(dtype_x, dtype_y, dtype_w, dtype_out):
    if (dtype_x == np.complex64) or (dtype_x == np.complex128):
        if (dtype_y == np.float64) or (dtype_y == np.float32):
            ys = "%s_fromreal(y[i])" % complex_dtype_to_name(dtype_x)
        else:
            ys = "y[i]"
        inner_map="%s_mul(%s_conj(x[i]), %s)" % (complex_dtype_to_name(dtype_x), complex_dtype_to_name(dtype_x), ys)
    else:
        inner_map="x[i]*y[i]"  
        
    if (dtype_w == np.float64) or (dtype_w == np.float32):
        inner_map = inner_map + "/w[i]"  
    else:
        inner_map = "%s_divide(%s, %s)" % (complex_dtype_to_name(dtype_x), inner_map, "w[i]")
               
    return ReductionKernel(mgr.state.context, dtype_out,
            neutral="0",
            arguments="__global const %(tp_x)s *x, __global const %(tp_y)s *y, __global const %(tp_w)s *w" % {
                "tp_x": dtype_to_ctype(dtype_x),
                "tp_y": dtype_to_ctype(dtype_y),
                "tp_w": dtype_to_ctype(dtype_w),
                },
            reduce_expr="a+b",
            map_expr=inner_map,
            name="weighted_inner")

def _get_dot_expr(dtype_out, dtype_a, dtype_b, conjugate_first,
        has_double_support, index_expr="i"):
    if dtype_b is None:
        if dtype_a is None:
            dtype_b = dtype_out
        else:
            dtype_b = dtype_a

    if dtype_out is None:
        from pyopencl.compyte.array import get_common_dtype
        dtype_out = get_common_dtype(
                dtype_a.type(0), dtype_b.type(0),
                has_double_support)

    a_real_dtype = dtype_a.type(0).real.dtype
    b_real_dtype = dtype_b.type(0).real.dtype
    out_real_dtype = dtype_out.type(0).real.dtype

    a_is_complex = dtype_a.kind == "c"
    b_is_complex = dtype_b.kind == "c"
    out_is_complex = dtype_out.kind == "c"

    from pyopencl.elementwise import complex_dtype_to_name

    if a_is_complex and b_is_complex:
        a = "a[%s]" % index_expr
        b = "b[%s]" % index_expr
        if dtype_a != dtype_out:
            a = "%s_cast(%s)" % (complex_dtype_to_name(dtype_out), a)
        if dtype_b != dtype_out:
            b = "%s_cast(%s)" % (complex_dtype_to_name(dtype_out), b)

        if conjugate_first and a_is_complex:
            a = "%s_conj(%s)" % (
                    complex_dtype_to_name(dtype_out), a)

        map_expr = "%s_mul(%s, %s)" % (
                complex_dtype_to_name(dtype_out), a, b)
    else:
        a = "a[%s]" % index_expr
        b = "b[%s]" % index_expr

        if out_is_complex:
            if a_is_complex and dtype_a != dtype_out:
                a = "%s_cast(%s)" % (complex_dtype_to_name(dtype_out), a)
            if b_is_complex and dtype_b != dtype_out:
                b = "%s_cast(%s)" % (complex_dtype_to_name(dtype_out), b)

            if not a_is_complex and a_real_dtype != out_real_dtype:
                a = "(%s) (%s)" % (dtype_to_ctype(out_real_dtype), a)
            if not b_is_complex and b_real_dtype != out_real_dtype:
                b = "(%s) (%s)" % (dtype_to_ctype(out_real_dtype), b)

        if conjugate_first and a_is_complex:
            a = "%s_conj(%s)" % (
                    complex_dtype_to_name(dtype_out), a)

        map_expr = "%s*%s" % (a, b)

    return map_expr, dtype_out, dtype_b

def get_divide_kernel(context, dtype_x, dtype_y, dtype_z):
    x_is_complex = dtype_x.kind == "c"
    y_is_complex = dtype_y.kind == "c"
    z_is_complex = dtype_z.kind == "c"

    x = "x[i]"
    y = "y[i]"

    if z_is_complex and dtype_x != dtype_y:
        if x_is_complex and dtype_x != dtype_z:
            x = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), x)
        if y_is_complex and dtype_y != dtype_z:
            y = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), y)

    if x_is_complex and y_is_complex:
        xoy = "%s_divide(%s, %s)" % (complex_dtype_to_name(dtype_z), x, y)
    elif not x_is_complex and y_is_complex:
        xoy = "%s_rdivide(%s, %s)" % (complex_dtype_to_name(dtype_z), x, y)
    elif x_is_complex and not y_is_complex:
        xoy = "%s_divider(%s, %s)" % (complex_dtype_to_name(dtype_z), x, y)
    else:
        xoy = "%s / %s" % (x, y)

    if z_is_complex:
        xoy = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), xoy)

    return get_elwise_kernel(context,
            "%(tp_z)s *z, %(tp_x)s *x, %(tp_y)s *y" % {
                "tp_x": dtype_to_ctype(dtype_x),
                "tp_y": dtype_to_ctype(dtype_y),
                "tp_z": dtype_to_ctype(dtype_z),
                },
            "z[i] = %s" % xoy,
            name="divide")

def get_multiply_kernel(context, dtype_x, dtype_y, dtype_z):
    x_is_complex = dtype_x.kind == "c"
    y_is_complex = dtype_y.kind == "c"

    x = "x[i]"
    y = "y[i]"

    if x_is_complex and dtype_x != dtype_z:
        x = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), x)
    if y_is_complex and dtype_y != dtype_z:
        y = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), y)

    if x_is_complex and y_is_complex:
        xy = "%s_mul(%s, %s)" % (complex_dtype_to_name(dtype_z), x, y)
    elif x_is_complex and not y_is_complex:
        xy = "%s_mulr(%s, %s)" % (complex_dtype_to_name(dtype_z), x, y)
    elif not x_is_complex and y_is_complex:
        xy = "%s_rmul(%s, %s)" % (complex_dtype_to_name(dtype_z), x, y)
    else:
        xy = "%s * %s" % (x, y)

    return get_elwise_kernel(context,
            "%(tp_z)s *z, %(tp_x)s *x, %(tp_y)s *y" % {
                "tp_x": dtype_to_ctype(dtype_x),
                "tp_y": dtype_to_ctype(dtype_y),
                "tp_z": dtype_to_ctype(dtype_z),
                },
            "z[i] = %s" % xy,
            name="multiply")

def get_axpbyz_kernel(context, dtype_x, dtype_y, dtype_z):
    ax = "a*x[i]"
    by = "b*y[i]"

    x_is_complex = dtype_x.kind == "c"
    y_is_complex = dtype_y.kind == "c"

    if x_is_complex:
        ax = "%s_mul(a, x[i])" % complex_dtype_to_name(dtype_x)

    if y_is_complex:
        by = "%s_mul(b, y[i])" % complex_dtype_to_name(dtype_y)

    if x_is_complex and not y_is_complex:
        by = "%s_fromreal(%s)" % (complex_dtype_to_name(dtype_x), by)

    if not x_is_complex and y_is_complex:
        ax = "%s_fromreal(%s)" % (complex_dtype_to_name(dtype_y), ax)

    if x_is_complex or y_is_complex:
        result = "{root}_add({root}_cast({ax}), {root}_cast({by}))".format(
            ax=ax, by=by, root=complex_dtype_to_name(dtype_z)
        )
    else:
        result = "%s + %s" % (ax, by)

    return get_elwise_kernel(
        context,
        "%(tp_z)s *z, %(tp_x)s a, %(tp_x)s *x, %(tp_y)s b, %(tp_y)s *y"
        % {"tp_x": dtype_to_ctype(dtype_x), "tp_y": dtype_to_ctype(dtype_y), "tp_z": dtype_to_ctype(dtype_z)},
        "z[i] = %s" % result,
        name="axpbyz",
    )

def get_take_put_kernel(context, dtype, idx_dtype, with_offsets, vec_count=1):
    ctx = {"idx_tp": dtype_to_ctype(idx_dtype), "tp": dtype_to_ctype(dtype)}

    args = (
        [VectorArg(dtype, "dest%d" % i) for i in range(vec_count)]
        + [
            VectorArg(idx_dtype, "gmem_dest_idx", with_offset=True),
            VectorArg(idx_dtype, "gmem_src_idx", with_offset=True),
        ]
        + [VectorArg(dtype, "src%d" % i, with_offset=True) for i in range(vec_count)]
        + [ScalarArg(idx_dtype, "offset%d" % i) for i in range(vec_count) if with_offsets]
    )

    if with_offsets:

        def get_copy_insn(i):
            return "dest%d[dest_idx] = " "src%d[src_idx+offset%d];" % (i, i, i)

    else:

        def get_copy_insn(i):
            return "dest%d[dest_idx] = " "src%d[src_idx];" % (i, i)

    body = ("%(idx_tp)s src_idx = gmem_src_idx[i];\n" "%(idx_tp)s dest_idx = gmem_dest_idx[i];\n" % ctx) + "\n".join(
        get_copy_insn(i) for i in range(vec_count)
    )

    return get_elwise_kernel(
        context, args, body, preamble=dtype_to_c_struct(context.devices[0], dtype), name="take_put"
    )

def get_axpbyz_kernel(context, dtype_x, dtype_y, dtype_z):
    ax = "a*x[i]"
    by = "b*y[i]"

    x_is_complex = dtype_x.kind == "c"
    y_is_complex = dtype_y.kind == "c"
    z_is_complex = dtype_z.kind == "c"

    if x_is_complex:
        ax = "%s_mul(a, x[i])" % complex_dtype_to_name(dtype_x)

    if y_is_complex:
        by = "%s_mul(b, y[i])" % complex_dtype_to_name(dtype_y)

    if x_is_complex and not y_is_complex:
        by = "%s_fromreal(%s)" % (complex_dtype_to_name(dtype_x), by)

    if not x_is_complex and y_is_complex:
        ax = "%s_fromreal(%s)" % (complex_dtype_to_name(dtype_y), ax)

    result = "%s + %s" % (ax, by)
    if z_is_complex:
        result = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), result)

    return get_elwise_kernel(context,
            "%(tp_z)s *z, %(tp_x)s a, %(tp_x)s *x, %(tp_y)s b, %(tp_y)s *y" % {
                "tp_x": dtype_to_ctype(dtype_x),
                "tp_y": dtype_to_ctype(dtype_y),
                "tp_z": dtype_to_ctype(dtype_z),
                },
            "z[i] = %s" % result,
            name="axpbyz")

def get_rdivide_elwise_kernel(context, dtype_x, dtype_y, dtype_z):
    # implements y / x!
    x_is_complex = dtype_x.kind == "c"
    y_is_complex = dtype_y.kind == "c"
    z_is_complex = dtype_z.kind == "c"

    x = "x[i]"
    y = "y"

    if z_is_complex and dtype_x != dtype_y:
        if x_is_complex and dtype_x != dtype_z:
            x = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), x)
        if y_is_complex and dtype_y != dtype_z:
            y = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), y)

    if x_is_complex and y_is_complex:
        yox = "%s_divide(%s, %s)" % (complex_dtype_to_name(dtype_z), y, x)
    elif not y_is_complex and x_is_complex:
        yox = "%s_rdivide(%s, %s)" % (complex_dtype_to_name(dtype_z), y, x)
    elif y_is_complex and not x_is_complex:
        yox = "%s_divider(%s, %s)" % (complex_dtype_to_name(dtype_z), y, x)
    else:
        yox = "%s / %s" % (y, x)

    return get_elwise_kernel(context,
            "%(tp_z)s *z, %(tp_x)s *x, %(tp_y)s y" % {
                "tp_x": dtype_to_ctype(dtype_x),
                "tp_y": dtype_to_ctype(dtype_y),
                "tp_z": dtype_to_ctype(dtype_z),
                },
            "z[i] = %s" % yox,
            name="divide_r")

def get_take_put_kernel(context, dtype, idx_dtype, with_offsets, vec_count=1):
    ctx = {
            "idx_tp": dtype_to_ctype(idx_dtype),
            "tp": dtype_to_ctype(dtype),
            }

    args = [
            VectorArg(dtype, "dest%d" % i)
                for i in range(vec_count)
            ] + [
            VectorArg(idx_dtype, "gmem_dest_idx"),
            VectorArg(idx_dtype, "gmem_src_idx"),
            ] + [
            VectorArg(dtype, "src%d" % i)
                for i in range(vec_count)
            ] + [
            ScalarArg(idx_dtype, "offset%d" % i)
                for i in range(vec_count) if with_offsets
            ]

    if with_offsets:
        def get_copy_insn(i):
            return ("dest%d[dest_idx] = "
                    "src%d[src_idx+offset%d];"
                    % (i, i, i))
    else:
        def get_copy_insn(i):
            return ("dest%d[dest_idx] = "
                    "src%d[src_idx];" % (i, i))

    body = (("%(idx_tp)s src_idx = gmem_src_idx[i];\n"
                "%(idx_tp)s dest_idx = gmem_dest_idx[i];\n" % ctx)
            + "\n".join(get_copy_insn(i) for i in range(vec_count)))

    return get_elwise_kernel(context, args, body, name="take_put")

def get_copy_kernel(context, dtype_dest, dtype_src):
    return get_elwise_kernel(context,
            "%(tp_dest)s *dest, %(tp_src)s *src" % {
                "tp_dest": dtype_to_ctype(dtype_dest),
                "tp_src": dtype_to_ctype(dtype_src),
                },
            "dest[i] = src[i]",
            name="copy")

def get_axpbz_kernel(context, dtype_a, dtype_x, dtype_b, dtype_z):
    a_is_complex = dtype_a.kind == "c"
    x_is_complex = dtype_x.kind == "c"
    b_is_complex = dtype_b.kind == "c"

    z_is_complex = dtype_z.kind == "c"

    ax = "a*x[i]"
    if x_is_complex:
        a = "a"
        x = "x[i]"

        if dtype_x != dtype_z:
            x = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), x)

        if a_is_complex:
            if dtype_a != dtype_z:
                a = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), a)

            ax = "%s_mul(%s, %s)" % (complex_dtype_to_name(dtype_z), a, x)
        else:
            ax = "%s_rmul(%s, %s)" % (complex_dtype_to_name(dtype_z), a, x)
    elif a_is_complex:
        a = "a"
        x = "x[i]"

        if dtype_a != dtype_z:
            a = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), a)
        ax = "%s_mulr(%s, %s)" % (complex_dtype_to_name(dtype_z), a, x)

    b = "b"
    if z_is_complex and not b_is_complex:
        b = "%s_fromreal(%s)" % (complex_dtype_to_name(dtype_z), b)

    if z_is_complex and not (a_is_complex or x_is_complex):
        ax = "%s_fromreal(%s)" % (complex_dtype_to_name(dtype_z), ax)

    if z_is_complex:
        ax = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), ax)
        b = "%s_cast(%s)" % (complex_dtype_to_name(dtype_z), b)

    if a_is_complex or x_is_complex or b_is_complex:
        expr = "{root}_add({ax}, {b})".format(
                ax=ax,
                b=b,
                root=complex_dtype_to_name(dtype_z))
    else:
        expr = "%s + %s" % (ax, b)

    return get_elwise_kernel(context,
            "%(tp_z)s *z, %(tp_a)s a, %(tp_x)s *x,%(tp_b)s b" % {
                "tp_a": dtype_to_ctype(dtype_a),
                "tp_x": dtype_to_ctype(dtype_x),
                "tp_b": dtype_to_ctype(dtype_b),
                "tp_z": dtype_to_ctype(dtype_z),
                },
            "z[i] = " + expr,
            name="axpb")

def get_multiply_kernel(context, dtype_x, dtype_y, dtype_z):
    return get_elwise_kernel(context,
            "%(tp_z)s *z, %(tp_x)s *x, %(tp_y)s *y" % {
                "tp_x": dtype_to_ctype(dtype_x),
                "tp_y": dtype_to_ctype(dtype_y),
                "tp_z": dtype_to_ctype(dtype_z),
                },
            "z[i] = x[i] * y[i]",
            name="multiply")

def get_pow_array_kernel(context, dtype_x, dtype_y, dtype_z):
    return get_elwise_kernel(context,
            "%(tp_z)s *z, %(tp_x)s *x, %(tp_y)s *y" % {
                "tp_x": dtype_to_ctype(dtype_x),
                "tp_y": dtype_to_ctype(dtype_y),
                "tp_z": dtype_to_ctype(dtype_z),
                },
            "z[i] = pow(x[i], y[i])",
            name="pow_method")

def get_axpbyz_kernel(context, dtype_x, dtype_y, dtype_z):
    return get_elwise_kernel(context,
            "%(tp_z)s *z, %(tp_x)s a, %(tp_x)s *x, %(tp_y)s b, %(tp_y)s *y" % {
                "tp_x": dtype_to_ctype(dtype_x),
                "tp_y": dtype_to_ctype(dtype_y),
                "tp_z": dtype_to_ctype(dtype_z),
                },
            "z[i] = a*x[i] + b*y[i]",
            name="axpbyz")

def get_divide_kernel(context, dtype_x, dtype_y, dtype_z):
    return get_elwise_kernel(context,
            "%(tp_z)s *z, %(tp_x)s *x, %(tp_y)s *y" % {
                "tp_x": dtype_to_ctype(dtype_x),
                "tp_y": dtype_to_ctype(dtype_y),
                "tp_z": dtype_to_ctype(dtype_z),
                },
            "z[i] = x[i] / y[i]",
            name="divide")

def get_unary_func_kernel(context, func_name, in_dtype, out_dtype=None):
    if out_dtype is None:
        out_dtype = in_dtype

    return get_elwise_kernel(context,
            "%(tp_out)s *z, %(tp_in)s *y" % {
                "tp_in": dtype_to_ctype(in_dtype),
                "tp_out": dtype_to_ctype(out_dtype),
                },
            "z[i] = %s(y[i])" % func_name,
            name="%s_kernel" % func_name)

def get_copy_kernel(context, dtype_dest, dtype_src):
    src = "src[i]"
    if dtype_dest.kind == "c" != dtype_src.kind:
        src = "%s_fromreal(%s)" % (complex_dtype_to_name(dtype_dest), src)

    return get_elwise_kernel(context,
            "%(tp_dest)s *dest, %(tp_src)s *src" % {
                "tp_dest": dtype_to_ctype(dtype_dest),
                "tp_src": dtype_to_ctype(dtype_src),
                },
            "dest[i] = %s" % src,
            name="copy")

def get_norm_kernel(dtype_x, dtype_out):
    if dtype_x == np.float32 or dtype_x == np.float64:
        op = "z[i] = x[i] * x[i]"
    if dtype_x == np.complex64 or dtype_x == np.complex128:
        op = "z[i] = x[i].x*x[i].x + x[i].y*x[i].y"
    return ElementwiseKernel(mgr.state.context, 
            "%(tp_x)s *x, %(tp_z)s *z" % {
                "tp_x": dtype_to_ctype(dtype_x),
                "tp_z": dtype_to_ctype(dtype_out),
                },
            op,
            "normsq")

def get_take_kernel(context, dtype, idx_dtype, vec_count=1):
    ctx = {"idx_tp": dtype_to_ctype(idx_dtype), "tp": dtype_to_ctype(dtype)}

    args = (
        [VectorArg(dtype, "dest" + str(i), with_offset=True) for i in range(vec_count)]
        + [VectorArg(dtype, "src" + str(i), with_offset=True) for i in range(vec_count)]
        + [VectorArg(idx_dtype, "idx", with_offset=True)]
    )
    body = ("%(idx_tp)s src_idx = idx[i];\n" % ctx) + "\n".join(
        "dest%d[i] = src%d[src_idx];" % (i, i) for i in range(vec_count)
    )

    return get_elwise_kernel(context, args, body, preamble=dtype_to_c_struct(context.devices[0], dtype), name="take")