import cupy

from cupy import _util
from cupyx.scipy.ndimage import _filters_core


def _get_sub_kernel(f):
    """
    Takes the "function" given to generic_filter and returns the "sub-kernel"
    that will be called, one of RawKernel or ReductionKernel.

    This supports:
     * cupy.RawKernel
       no checks are possible
     * cupy.ReductionKernel
       checks that there is a single input and output
    """
    if isinstance(f, cupy.RawKernel):
        # We will assume that it has the correct API
        return f
    elif isinstance(f, cupy.ReductionKernel):
        if f.nin != 1 or f.nout != 1:
            raise TypeError('ReductionKernel must have 1 input and output')
        return f
    elif isinstance(f, cupy.ElementwiseKernel):
        # special error message for ElementwiseKernels
        raise TypeError('only ReductionKernel allowed (not ElementwiseKernel)')
    else:
        raise TypeError('bad function type')


@_util.memoize(for_each_device=True)
def _get_generic_filter_red(rk, in_dtype, out_dtype, filter_size, mode,
                            wshape, offsets, cval, int_type):
    """Generic filter implementation based on a reduction kernel."""
    # Get the temporary output c type
    in_param, out_param = rk.in_params[0], rk.out_params[0]
    out_ctype = out_param.ctype
    if out_param.dtype is None:  # resolve template
        out_ctype = cupy._core._scalar.get_typename(
            in_dtype if out_param.ctype == in_param.ctype else out_dtype)

    # Get code chunks
    setup = '''
    int iv = 0;
    X values[{size}];
    CArray<X, 1, true, true> sub_in(values, {{{size}}});
    {out_ctype} val_out;
    CArray<{out_ctype}, 1, true, true> sub_out(&val_out, {{1}});
    '''.format(size=filter_size, out_ctype=out_ctype)

    sub_call = '''reduction_kernel::{}(sub_in, sub_out);
    y = cast<Y>(val_out);'''.format(rk.name)

    sub_kernel = _reduction_kernel_code(rk, filter_size, out_dtype, in_dtype)

    # Get the final kernel
    return _filters_core._generate_nd_kernel(
        'generic_{}_{}'.format(filter_size, rk.name),
        setup, 'values[iv++] = {value};', sub_call,
        mode, wshape, int_type, offsets, cval, preamble=sub_kernel,
        options=getattr(rk, 'options', ()))


def _reduction_kernel_code(rk, filter_size, out_dtype, in_dtype):
    # NOTE: differences from the code generated for real reduction kernels:
    #  * input is always 1D and always less than 2^31 elements
    #  * output is always 1D with a single element
    #  * never across threads (no _block_stride, _sdata, _sdata_raw, _REDUCE,
    #       _tid, _J, _i, _i_base, _j_offset, _J_offset, _j_stride, _J_stride)
    # Also, the code is moved into a namespace so that clashes are minimized
    # between the typedefs for the "template" variables.

    # figure out the types
    types = {}
    in_param, out_param = rk.in_params[0], rk.out_params[0]
    in_ctype = _get_type_info(in_param, in_dtype, types)
    out_ctype = _get_type_info(out_param, out_dtype, types)
    types = '\n'.join('typedef {} {};'.format(typ, name)
                      for name, typ in types.items())

    return '''namespace reduction_kernel {{
{type_preamble}
{preamble}
__device__
void {name}({in_const} CArray<{in_ctype}, 1, true, true>& _raw_{in_name},
            CArray<{out_ctype}, 1, true, true>& _raw_{out_name}) {{
    // these are just provided so if they are available for the RK
    CIndexer<1> _in_ind({{{size}}});
    CIndexer<0> _out_ind;

    #define REDUCE(a, b) ({reduce_expr})
    #define POST_MAP(a) ({post_map_expr})
    typedef {reduce_type} _type_reduce;
    _type_reduce _s = _type_reduce({identity});
    for (int _j = 0; _j < {size}; ++_j) {{
        _in_ind.set(_j);
        {in_const} {in_ctype}& {in_name} = _raw_{in_name}[_j];
        _type_reduce _a = static_cast<_type_reduce>({pre_map_expr});
        _s = REDUCE(_s, _a);
    }}
    _out_ind.set(0);
    {out_ctype} &{out_name} = _raw_{out_name}[0];
    POST_MAP(_s);
    #undef REDUCE
    #undef POST_MAP
}}
}}'''.format(
        name=rk.name, type_preamble=types, preamble=rk.preamble,
        in_const='const' if in_param.is_const else '',
        in_ctype=in_ctype, in_name=in_param.name,
        out_ctype=out_ctype, out_name=out_param.name,

        pre_map_expr=rk.map_expr,
        identity='' if rk.identity is None else rk.identity,
        size=filter_size,
        reduce_type=rk.reduce_type, reduce_expr=rk.reduce_expr,
        post_map_expr=rk.post_map_expr,
    )


def _get_type_info(param, dtype, types):
    if param.dtype is not None:
        return param.ctype
    # Template type -> map to actual output type
    ctype = cupy._core._scalar.get_typename(dtype)
    types.setdefault(param.ctype, ctype)
    return ctype


@_util.memoize(for_each_device=True)
def _get_generic_filter_raw(rk, filter_size, mode, wshape, offsets, cval,
                            int_type):
    """Generic filter implementation based on a raw kernel."""
    setup = '''
    int iv = 0;
    double values[{}];
    double val_out;'''.format(filter_size)

    sub_call = '''raw_kernel::{}(values, {}, &val_out);
    y = cast<Y>(val_out);'''.format(rk.name, filter_size)

    return _filters_core._generate_nd_kernel(
        'generic_{}_{}'.format(filter_size, rk.name),
        setup, 'values[iv++] = cast<double>({value});', sub_call,
        mode, wshape, int_type, offsets, cval,
        preamble='namespace raw_kernel {{\n{}\n}}'.format(
            # Users can test RawKernel independently, but when passed to here
            # it must be used as a device function here. In fact, RawKernel
            # wouldn't compile if code only contains device functions, so this
            # is necessary.
            rk.code.replace('__global__', '__device__')),
        options=rk.options)


@_util.memoize(for_each_device=True)
def _get_generic_filter1d(rk, length, n_lines, filter_size, origin, mode, cval,
                          in_ctype, out_ctype, int_type):
    """
    The generic 1d filter is different than all other filters and thus is the
    only filter that doesn't use _generate_nd_kernel() and has a completely
    custom raw kernel.
    """
    in_length = length + filter_size - 1
    start = filter_size // 2 + origin
    end = start + length

    if mode == 'constant':
        boundary, boundary_early = '', '''
        for (idx_t j = 0; j < {start}; ++j) {{ input_line[j] = {cval}; }}
        for (idx_t j = {end}; j<{in_length}; ++j) {{ input_line[j] = {cval}; }}
        '''.format(start=start, end=end, in_length=in_length, cval=cval)
    else: