"""
Ops for downsampling images.
Planned:
Pool, DownsampleAvg, DownsampleSoftmax.
"""
import itertools
import warnings

import numpy as np

import aesara.tensor.basic as at
import aesara.tensor.math as tm
from aesara.gradient import DisconnectedType
from aesara.graph.basic import Apply, Constant, Variable
from aesara.graph.params_type import ParamsType
from aesara.graph.type import EnumList
from aesara.graph.utils import MethodNotDefined
from aesara.link.c.op import OpenMPOp
from aesara.scalar import bool as bool_t
from aesara.tensor.type import TensorType, int_dtypes


def max_pool_2d_same_size(input, patch_size):
    """
    Takes as input a 4-D tensor. It sets all non maximum values
    of non-overlapping patches of size (patch_size[0],patch_size[1]) to zero,
    keeping only the maximum values. The output has the same dimensions as
    the input.

    Parameters
    ----------
    input : 4-D aesara tensor of input images
        Input images. Max pooling will be done over the 2 last dimensions.
    patch_size : tuple of length 2 or aesara vector of ints of size 2.
        Size of the patch (patch height, patch width).
        (2,2) will retain only one non-zero value per patch of 4 values.

    """
    output = Pool(True)(input, patch_size)
    outs = MaxPoolGrad(True)(input, output, output, patch_size)
    return outs


def pool_2d(
    input,
    ws=None,
    ignore_border=None,
    stride=None,
    pad=(0, 0),
    mode="max",
    ds=None,
    st=None,
    padding=None,
):
    """Downscale the input by a specified factor

    Takes as input a N-D tensor, where N >= 2. It downscales the input image by
    the specified factor, by keeping only the maximum value of non-overlapping
    patches of size (ws[0],ws[1])

    Parameters
    ----------
    input : N-D aesara tensor of input images
        Input images. Max pooling will be done over the 2 last dimensions.
    ws : tuple of length 2 or aesara vector of ints of size 2.
        Factor by which to downscale (vertical ws, horizontal ws).
        (2,2) will halve the image in each dimension.
    ignore_border : bool (default None, will print a warning and set to False)
        When True, (5,5) input with ws=(2,2) will generate a (2,2) output.
        (3,3) otherwise.
    stride : tuple of two ints or aesara vector of ints of size 2.
        Stride size, which is the number of shifts over rows/cols to get the
        next pool region. If stride is None, it is considered equal to ws
        (no overlap on pooling regions), eg: stride=(1,1) will shifts over
        one row and one col for every iteration.
    pad : tuple of two ints or aesara vector of ints of size 2.
        (pad_h, pad_w), pad zeros to extend beyond four borders of the
        images, pad_h is the size of the top and bottom margins, and
        pad_w is the size of the left and right margins.
    mode : {'max', 'sum', 'average_inc_pad', 'average_exc_pad'}
        Operation executed on each window. `max` and `sum` always exclude
        the padding in the computation. `average` gives you the choice to
        include or exclude it.
    ds
        *deprecated*, use parameter ws instead.
    st
        *deprecated*, use parameter stride instead.
    padding
        *deprecated*, use parameter pad instead.

    """
    # check for deprecated parameter names
    if ds is not None:
        if ws is not None:
            raise ValueError(
                "You can't provide a tuple value to both 'ws' and 'ds'."
                " Please provide a value only to 'ws'."
            )
        else:
            warnings.warn(
                "The 'ds' parameter is not going to exist"
                " anymore as it is going to be replaced by the parameter"
                " 'ws'.",
                category=DeprecationWarning,
                stacklevel=2,
            )
            ws = ds
    elif ds is None and ws is None:
        raise ValueError("You must provide a tuple value for the window size.")

    if st is not None:
        if stride is not None:
            raise ValueError(
                "You can't provide a tuple value to both 'st and 'stride'."
                " Please provide a value only to 'stride'."
            )
        else:
            warnings.warn(
                "The 'st' parameter is not going to exist"
                " anymore as it is going to be replaced by the parameter"
                " 'stride'.",
                category=DeprecationWarning,
                stacklevel=2,
            )
            stride = st

    if padding is not None:
        if pad not in {None, (0, 0)}:
            raise ValueError(
                "You can't provide a tuple value to both 'padding' and pad."
                "  Please provide a value only to pad."
            )
        else:
            warnings.warn(
                "The 'padding' parameter is not going to exist"
                " anymore as it is going to be replaced by the parameter"
                " 'pad'.",
                category=DeprecationWarning,
                stacklevel=2,
            )
            pad = padding

    if input.ndim < 2:
        raise NotImplementedError("pool_2d requires a dimension >= 2")
    if ignore_border is None:
        warnings.warn(
            "pool_2d() will have the parameter ignore_border"
            " default value changed to True (currently"
            " False). To have consistent behavior with all Aesara"
            " version, explicitly add the parameter ignore_border=True."
            " On the GPU, using ignore_border=True is needed to use cuDNN."
            " When using ignore_border=False and not using cuDNN, the only"
            " GPU combination supported is when"
            " `ws == stride and pad == (0, 0) and mode == 'max'`."
            " Otherwise, the convolution will be executed on CPU.",
            category=DeprecationWarning,
            stacklevel=2,
        )
        ignore_border = False
    op = Pool(ignore_border, ndim=2, mode=mode)
    output = op(input, ws, stride, pad)
    return output


def pool_3d(
    input,
    ws=None,
    ignore_border=None,
    stride=None,
    pad=(0, 0, 0),
    mode="max",
    ds=None,
    st=None,
    padding=None,
):
    """Downscale the input by a specified factor

    Takes as input a N-D tensor, where N >= 3. It downscales the input image by
    the specified factor, by keeping only the maximum value of non-overlapping
    patches of size (ws[0],ws[1],ws[2])

    Parameters
    ----------
    input : N-D aesara tensor of input images
        Input images. Max pooling will be done over the 3 last dimensions.
    ws : tuple of length 3 or aesara vector of ints of size 3
        Factor by which to downscale (vertical ws, horizontal ws, depth ws).
        (2,2,2) will halve the image in each dimension.
    ignore_border : bool (default None, will print a warning and set to False)
        When True, (5,5,5) input with ws=(2,2,2) will generate a (2,2,2) output.
        (3,3,3) otherwise.
    st : tuple of three ints or aesara vector of ints of size 3
        Stride size, which is the number of shifts over rows/cols/slices to get
        the next pool region. If st is None, it is considered equal to ws
        (no overlap on pooling regions).
    pad : tuple of two ints or aesara vector of ints of size 3
        (pad_h, pad_w, pad_d), pad zeros to extend beyond six borders of the
        images, pad_h is the size of the top and bottom margins,
        pad_w is the size of the left and right margins, and pad_d is the size
        of the front and back margins
    mode : {'max', 'sum', 'average_inc_pad', 'average_exc_pad'}
        Operation executed on each window. `max` and `sum` always exclude
        the padding in the computation. `average` gives you the choice to
        include or exclude it.
    ds
        *deprecated*, use parameter ws instead.
    st
        *deprecated*, use parameter st instead.
    padding
        *deprecated*, use parameter pad instead.

    """
    # check for deprecated parameter names
    if ds is not None:
        if ws is not None:
            raise ValueError(
                "You can't provide a tuple value to both 'ws' and 'ds'."
                " Please provide a value only to 'ws'."
            )
        else:
            warnings.warn(
                "The 'ds' parameter is not going to exist"
                " anymore as it is going to be replaced by the parameter"
                " 'ws'.",
                category=DeprecationWarning,
                stacklevel=2,
            )
            ws = ds
    elif ds is None and ws is None:
        raise ValueError("You must provide a tuple value for the window size.")

    if st is not None:
        if stride is not None:
            raise ValueError(
                "You can't provide a tuple value to both 'st and 'stride'."
                " Please provide a value only to 'stride'."
            )
        else:
            warnings.warn(
                "The 'st' parameter is not going to exist"
                " anymore as it is going to be replaced by the parameter"
                " 'stride'.",
                category=DeprecationWarning,
                stacklevel=2,
            )
            stride = st

    if padding is not None:
        if pad not in {None, (0, 0, 0)}:
            raise ValueError(
                "You can't provide a tuple value to both 'padding' and pad."
                "  Please provide a value only to pad."
            )
        else:
            warnings.warn(
                "The 'padding' parameter is not going to exist"
                " anymore as it is going to be replaced by the parameter"
                " 'pad'.",
                category=DeprecationWarning,
                stacklevel=2,
            )
            pad = padding

    if input.ndim < 3:
        raise NotImplementedError("pool_3d requires a dimension >= 3")
    if ignore_border is None:
        warnings.warn(
            "pool_3d() will have the parameter ignore_border"
            " default value changed to True (currently"
            " False). To have consistent behavior with all Aesara"
            " version, explicitly add the parameter ignore_border=True."
            " On the GPU, using ignore_border=True is needed to use cuDNN."
            " When using ignore_border=False and not using cuDNN, the only"
            " GPU combination supported is when"
            " `ws == stride and pad == (0, 0, 0) and mode == 'max'`."
            " Otherwise, the convolution will be executed on CPU.",
            category=DeprecationWarning,
            stacklevel=2,
        )
        ignore_border = False
    op = Pool(ignore_border, ndim=3, mode=mode)
    output = op(input, ws, stride, pad)
    return output


# NB: This enum type is currently used in gpuarray/pool.py.
# It may be used later as op param in this current file.
# Enum name and constants names are inspired from cuDNN type `cudnnPoolingMode_t`
# (cf. `aesara/gpuarray/cudnn_defs.py`).
PoolingMode_t = EnumList(
    ("POOLING_MAX", "max"),
    ("POOLING_SUM", "sum"),
    ("POOLING_AVERAGE_COUNT_INCLUDE_PADDING", "average_inc_pad"),
    ("POOLING_AVERAGE_COUNT_EXCLUDE_PADDING", "average_exc_pad"),
)


class Pool(OpenMPOp):
    """
    sum or average over different patches.

    Parameters
    ----------
    ws : list or tuple of N ints
        Downsample factor over rows, columns etc.
        ws indicates the size of the pooling region.
    ignore_border : bool
        If ws doesn't divide imgshape, do we include an extra row/col/slice
        of partial downsampling (False) or ignore it (True).
    stride : list or tuple of N ints or None
        Stride size, which is the number of shifts over rows/cols/slices to get the
        next pool region. If stride is None, it is considered equal to ws
        (no overlap on pooling regions).
    pad : tuple of N ints or None
        For each downsampling dimension, this specifies the number of zeros to
        add as padding on both sides. For 2D and (pad_h, pad_w), pad_h specifies the
        size of the top and bottom margins, pad_w specifies the size of the left and
        right margins. No padding is added if pad is None.
    mode : {'max', 'sum', 'average_inc_pad', 'average_exc_pad'}
        ('average_inc_pad' excludes the padding from the count,
        'average_exc_pad' include it)
    ndim : int
        The number of pooling dimensions N.
        The default is 2.
    ds
        *deprecated*, use parameter ws instead.
    st
        *deprecated*, use parameter st instead.
    padding
        *deprecated*, use parameter pad instead.


    """

    __props__ = ("ignore_border", "mode", "ndim")
    params_type = ParamsType(
        ignore_border=bool_t,
    )

    @staticmethod
    def out_shape(
        imgshape,
        ws=None,
        ignore_border=False,
        stride=None,
        pad=None,
        ndim=2,
        ds=None,
        st=None,
        padding=None,
    ):
        """
        Return the shape of the output from this op, for input of given
        shape and flags.

        Parameters
        ----------
        imgshape : tuple, list, or similar of integer or scalar Aesara variable
            The shape of a tensor of images. The last N elements are
            interpreted as the number of rows, and the number of cols.
        ws : list or tuple of N ints
            Downsample factor over rows and column.
            ws indicates the pool region size.
        ignore_border : bool
            If ws doesn't divide imgshape, do we include an extra row/col/slice
            of partial downsampling (False) or ignore it (True).
        stride : list or tuple of N ints or None
            Stride size, which is the number of shifts over rows/cols/slices to get the
            next pool region. If stride is None, it is considered equal to ws
            (no overlap on pooling regions).
        pad : tuple of N ints or None
            For each downsampling dimension, this specifies the number of zeros to
            add as padding on both sides. For 2D and (pad_h, pad_w), pad_h specifies the
            size of the top and bottom margins, pad_w specifies the size of the left and
            right margins. No padding is added if pad is None.
        ndim : int
            The number of pooling dimensions N.
            The default is 2.
        ds
            *deprecated*, use parameter ws instead.
        st
            *deprecated*, use parameter st instead.
        padding
            *deprecated*, use parameter pad instead.

        Returns
        -------
        list
            The shape of the output from this op, for input of given shape.
            This will have the same length as imgshape, but with last N
            elements reduced as per the downsampling & ignore_border flags.

        """
        # check for deprecated parameter names
        if ds is not None:
            if ws is not None:
                raise ValueError(
                    "You can't provide a tuple value to both 'ws' and 'ds'."
                    " Please provide a value only to 'ws'."
                )
            else:
                warnings.warn(
                    "The 'ds' parameter is not going to exist"
                    " anymore as it is going to be replaced by the parameter"
                    " 'ws'.",
                    category=DeprecationWarning,
                    stacklevel=2,
                )
                ws = ds
        elif ds is None and ws is None:
            raise ValueError("You must provide a tuple value for the window size.")

        if st is not None:
            if stride is not None:
                raise ValueError(
                    "You can't provide a tuple value to both 'st and 'stride'."
                    " Please provide a value only to 'stride'."
                )
            else:
                warnings.warn(
                    "The 'st' parameter is not going to exist"
                    " anymore as it is going to be replaced by the parameter"
                    " 'stride'.",
                    category=DeprecationWarning,
                    stacklevel=2,
                )
                stride = st

        if padding is not None:
            zero_pad = (0,) * ndim
            if pad not in {None, zero_pad}:
                raise ValueError(
                    "You can't provide a tuple value to both 'padding' and pad."
                    "  Please provide a value only to pad."
                )
            else:
                warnings.warn(
                    "The 'padding' parameter is not going to"
                    " exist anymore as it is going to be replaced by the"
                    " parameter 'pad'.",
                    category=DeprecationWarning,
                    stacklevel=2,
                )
                pad = padding

        if ndim is None:
            ndim = 2
        assert ndim > 0
        if len(imgshape) < ndim:
            raise TypeError(f"imgshape must have at least {ndim} dimensions")

        if stride is None:
            stride = ws
        if pad is None:
            pad = (0,) * ndim
        patch_shape = tuple(
            at.extract_constant(imgshape[-ndim + i]) + pad[i] * 2 for i in range(ndim)
        )

        def compute_out(v, downsample, stride):
            if ignore_border:
                if downsample == stride:
                    return v // stride
                else:
                    out = (v - downsample) // stride + 1
                    if isinstance(out, Variable):
                        return tm.maximum(out, 0)
                    else:
                        return np.maximum(out, 0)
            else:
                if isinstance(v, Variable):
                    return at.switch(
                        tm.ge(stride, downsample),
                        (v - 1) // stride + 1,
                        tm.maximum(0, (v - 1 - downsample) // stride + 1) + 1,
                    )
                elif stride >= downsample:
                    return (v - 1) // stride + 1
                else:
                    return max(0, (v - 1 - downsample + stride) // stride) + 1

        out_shape = [compute_out(patch_shape[i], ws[i], stride[i]) for i in range(ndim)]

        rval = list(imgshape[:-ndim]) + out_shape
        return rval

    def __init__(self, ignore_border=False, mode="max", ndim=2, openmp=None):
        super().__init__(openmp=openmp)
        self.ndim = ndim
        self.ignore_border = ignore_border
        if mode == "max_deterministic":
            # It seems max pool algo is already deterministic in CPU.
            mode = "max"
        if mode not in ("max", "average_inc_pad", "average_exc_pad", "sum"):
            raise ValueError(
                "Pool mode parameter only support 'max', 'sum',"
                f" 'average_inc_pad' and 'average_exc_pad'. Got {mode}"
            )
        self.mode = mode

    def prepare_node(self, node, storage_map, compute_map, impl):
        if len(node.inputs) == 1:
            # Old interface
            self.ndim = len(node.op.ds)
            self.mode = node.op.mode
            ws = at.constant(node.op.ds)
            st = at.constant(node.op.st)
            pad = at.constant(node.op.padding)
            node.inputs.append(ws)
            node.inputs.append(st)
            node.inputs.append(pad)
            if isinstance(ws, Constant):
                storage_map[ws] = [ws.data]
                compute_map[ws] = [True]
            else:
                storage_map[ws] = [None]
                compute_map[ws] = [False]
            if isinstance(st, Constant):
                storage_map[st] = [st.data]
                compute_map[st] = [True]
            else:
                storage_map[st] = [None]
                compute_map[st] = [False]
            if isinstance(pad, Constant):
                storage_map[pad] = [pad.data]
                compute_map[pad] = [True]
            else:
                storage_map[pad] = [None]
                compute_map[pad] = [False]

    def make_node(self, x, ws, stride=None, pad=None):
        # TODO: consider restricting the dtype?
        x = at.as_tensor_variable(x)
        nd = self.ndim
        if stride is None:
            stride = ws
        if pad is None:
            pad = (0,) * nd
        elif isinstance(pad, (tuple, list)):
            if max(pad) != 0 and not self.ignore_border:
                raise NotImplementedError("padding works only with ignore_border=True")
            if isinstance(ws, (tuple, list)):
                if any(pad[i] >= ws[i] for i in range(nd)):
                    raise NotImplementedError("padding must be smaller than strides")
        ws = at.as_tensor_variable(ws)
        stride = at.as_tensor_variable(stride)
        pad = at.as_tensor_variable(pad)
        assert ws.ndim == 1
        assert stride.ndim == 1
        assert pad.ndim == 1
        if x.type.ndim < nd:
            raise TypeError()
        if ws.dtype not in int_dtypes:
            raise TypeError("Pool downsample parameters must be ints.")
        if stride.dtype not in int_dtypes:
            raise TypeError("Stride parameters must be ints.")
        if pad.dtype not in int_dtypes:
            raise TypeError("Padding parameters must be ints.")
        # If the input shape are broadcastable we can have 0 in the output shape
        broad = x.broadcastable[:-nd] + (False,) * nd
        out = TensorType(x.dtype, broad)
        return Apply(self, [x, ws, stride, pad], [out()])

    def perform(self, node, inp, out, params):
        x, ws, stride, pad = inp
        (z,) = out
        nd = self.ndim
        assert ws.shape == stride.shape == pad.shape == (nd,)
        if len(x.shape) < nd:
            raise NotImplementedError(
                f"Pool requires input with {nd} or more dimensions"
            )
        z_shape = self.out_shape(x.shape, ws, params.ignore_border, stride, pad, nd)
        if not params.ignore_border:
            assert all(z > 0 for z in z_shape[-nd:])
        if (z[0] is None) or (z[0].shape != z_shape):
            z[0] = np.empty(z_shape, dtype=x.dtype)
        zz = z[0]
        # size of pooling output
        pool_out_shp = zz.shape[-nd:]
        img_shp = tuple(x.shape[-nd + i] + 2 * pad[i] for i in range(nd))
        inc_pad = self.mode == "average_inc_pad"

        # pad the image
        if max(pad) != 0:
            y = np.zeros(x.shape[:-nd] + img_shp, dtype=x.dtype)
            y[
                (slice(None),) * (len(x.shape) - nd)
                + tuple(slice(pad[i], img_shp[i] - pad[i]) for i in range(nd))
            ] = x
        else:
            y = x
        func = np.max
        if self.mode == "sum":
            func = np.sum
        elif self.mode != "max":
            func = np.average

        # precompute the region boundaries for each dimension
        region_slices = [[] for i in range(nd)]
        for i in range(nd):
            for j in range(pool_out_shp[i]):
                start = j * stride[i]
                end = min(start + ws[i], img_shp[i])
                if not inc_pad:
                    start = max(start, pad[i])
                    end = min(end, img_shp[i] - pad[i])
                region_slices[i].append(slice(start, end))

        # iterate over non-pooling dimensions
        for k in np.ndindex(*x.shape[:-nd]):
            zzk = zz[k]
            yk = y[k]
            # iterate over pooling regions
            for r in np.ndindex(*pool_out_shp):
                zzk[r] = func(yk[[region_slices[i][r[i]] for i in range(nd)]])

    def infer_shape(self, fgraph, node, in_shapes):
        ws, stride, pad = [node.inputs[1], node.inputs[2], node.inputs[3]]
        shp = self.out_shape(
            in_shapes[0], ws, self.ignore_border, stride, pad, self.ndim
        )
        return [shp]

    def L_op(self, inputs, outputs, grads):
        x, ws, stride, pad = inputs
        (gz,) = grads
        disc = [DisconnectedType()() for i in inputs[1:]]
        if self.mode == "max":
            return [
                MaxPoolGrad(ndim=self.ndim, ignore_border=self.ignore_border)(
                    x, outputs[0], gz, ws=ws, stride=stride, pad=pad
                )
            ] + disc
        else:
            return [
                AveragePoolGrad(
                    ndim=self.ndim, ignore_border=self.ignore_border, mode=self.mode
                )(x, gz, ws=ws, stride=stride, pad=pad)
            ] + disc

    def connection_pattern(self, node):
        return [[1], [0], [0], [0]]

    def R_op(self, inputs, eval_points):
        if self.mode != "max":
            # Rop for average or sum is simply pooling evaluated at eval point
            eval_inputs = [eval_points[0]] + inputs[1:]
            return [self(*eval_inputs)]

        # R_op can receive None as eval_points.
        # That mean there is no diferientiable path through that input
        # If this imply that you cannot compute some outputs,
        # return None for those.
        if eval_points[0] is None:
            return [None]
        z = self(*inputs)
        x, ws, stride, pad = inputs
        return [
            DownsampleFactorMaxGradGrad(self.ignore_border, self.mode, self.ndim)(
                x, z, eval_points[0], ws, stride, pad
            )
        ]

    def c_headers(self, **kwargs):
        headers = ["<algorithm>"]
        headers += super().c_headers(**kwargs)
        return headers

    def c_code(self, node, name, inp, out, sub):
        if self.mode not in ("max", "sum", "average_exc_pad", "average_inc_pad"):
            raise MethodNotDefined()
        x, ws, stride, pad = inp
        (z,) = out
        nd = self.ndim
        total_ndim = node.inputs[0].ndim
        non_pool_ndim = total_ndim - nd
        fail = sub["fail"]
        params = sub["params"]
        if self.openmp:
            # run in parallel over each pooling block
            omp_parallel = "#pragma omp parallel for private(r_st, r_end, r_idx, i_idx, o_idx, collector) schedule(static)"
        else:
            omp_parallel = ""
        ccode = """
        int typenum = PyArray_ObjectType((PyObject*)%(x)s, 0);
        if(PyArray_NDIM(%(x)s)!=%(total_ndim)s)
        {
            PyErr_SetString(PyExc_ValueError, "x must be a %(total_ndim)sD ndarray");
            %(fail)s;
        }
        if(PyArray_DIM(%(ws)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "ws must be a vector of size %(nd)s");
            %(fail)s;
        }
        if(PyArray_DIM(%(stride)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "stride must be a vector of size %(nd)s");
            %(fail)s;
        }
        if(PyArray_DIM(%(pad)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "pad must be a vector of size %(nd)s");
            %(fail)s;
        }
        npy_intp z[%(nd)s]; // shape of the output
        npy_intp r[%(nd)s]; // shape of the padded_input
        npy_intp ws[%(nd)s];
        npy_intp st[%(nd)s];
        npy_intp pd[%(nd)s];
        int nonzero_padding;
        nonzero_padding = 0;
        for (int i=0; i<%(nd)s; i++)
        {
            ws[i] = *((dtype_%(ws)s*)PyArray_GETPTR1(%(ws)s, i));
            st[i] = *((dtype_%(stride)s*)PyArray_GETPTR1(%(stride)s, i));
            pd[i] = *((dtype_%(pad)s*)PyArray_GETPTR1(%(pad)s, i));
            r[i] = PyArray_DIMS(%(x)s)[%(non_pool_ndim)s + i] + 2 * pd[i];
            if (pd[i]>0)
                nonzero_padding = 1;
        }
        if (!%(params)s->ignore_border && nonzero_padding)
        {
            PyErr_SetString(PyExc_ValueError,
              "padding must be zero when ignore border is False");
            %(fail)s;
        }
        if (%(params)s->ignore_border)
        {
            for (int i=0; i<%(nd)s; i++)
            {
                // '/' in C is different from '/' in python
                if (r[i] - ws[i] < 0)
                {
                  z[i] = 0;
                }
                else
                {
                  z[i] = (r[i] - ws[i]) / st[i] + 1;
                }
            }
        }
        else
        {
            for (int i=0; i<%(nd)s; i++)
            {
                // decide how many rows/cols the output has
                if (st[i] >= ws[i])
                {
                    z[i] = (r[i] - 1) / st[i] + 1;
                }
                else
                {
                    z[i] = std::max((npy_intp)0, (r[i] - 1 - ws[i] + st[i]) / st[i]) + 1;
                }
                assert(z[i] > 0);
            }
        }
        // memory allocation of z if necessary
        int mem_nec;
        mem_nec = 0;
        if ((!%(z)s) || *PyArray_DIMS(%(z)s)!=%(total_ndim)s)
        {
            mem_nec = 1;
        }
        if (!mem_nec)
        {
            for (int i=0; i<%(non_pool_ndim)s; i++)
            {
                if (PyArray_DIMS(%(z)s)[i] != PyArray_DIMS(%(x)s)[i])
                {
                    mem_nec = 1;
                    break;
                }
            }
        }
        if (!mem_nec)
        {
            for (int i=0; i<%(nd)s; i++)
            {
                if (PyArray_DIMS(%(z)s)[%(non_pool_ndim)s + i] != z[i])
                {
                    mem_nec = 1;
                    break;
                }
            }
        }
        if (mem_nec)
        {
          if (%(z)s) Py_XDECREF(%(z)s);
          npy_intp dims[%(total_ndim)s];
          for (int i=0; i<%(non_pool_ndim)s; i++)
          {
              dims[i] = PyArray_DIMS(%(x)s)[i];
          }
          for (int i=0; i<%(nd)s; i++)
          {
              dims[%(non_pool_ndim)s + i] = z[i];
          }
          //TODO: zeros not necessary
          %(z)s = (PyArrayObject*) PyArray_ZEROS(%(total_ndim)s, dims, typenum,0);
        }
        // initialize temp var for the value in a region
        dtype_%(x)s collector;
        npy_intp z_prod;
        // do not run if any z[i] is zero
        z_prod = 1;
        for (int i=0; i<%(nd)s; i++)
        {
            z_prod *= z[i];
        }
        if (z_prod)
        {
            // will be used to hold start and end index of a region
            npy_intp r_st[%(nd)s];
            npy_intp r_end[%(nd)s];
            // index for iterating over the pooling regions
            npy_intp r_idx[%(nd)s];
            // placeholder for PyArray indexing (output)
            npy_intp o_idx[%(total_ndim)s];
            // placeholder for PyArray indexing (input)
            npy_intp i_idx[%(total_ndim)s];
            // loop over non-pooling dimensions
            npy_intp non_pooling_prod = 1;
            for (int i=0; i<%(non_pool_ndim)s; i++)
            {
                non_pooling_prod *= PyArray_DIMS(%(x)s)[i];
            }
            %(omp_parallel)s
            // first loop over non-pooling dimensions
            for (npy_intp t=0; t<non_pooling_prod; t++)
            {
                // compute the non-pooling index in each dimension
                if (%(non_pool_ndim)s!=0)
                {
                    o_idx[0] = t;
                    i_idx[0] = t;
                    for (int i=1; i<%(non_pool_ndim)s; i++)
                    {
                        o_idx[i] = o_idx[i - 1] / PyArray_DIMS(%(x)s)[i - 1];
                        o_idx[i - 1] = o_idx[i - 1] %% PyArray_DIMS(%(x)s)[i - 1];
                        i_idx[i] = o_idx[i];
                        i_idx[i - 1] = o_idx[i - 1];
                    }
                }

                // then loop over each region in each pooling dimension
        """

        for i in range(nd):
            ccode += """
                for (r_idx[%(i)s]=0; r_idx[%(i)s] < z[%(i)s]; r_idx[%(i)s]++) {
                  r_st[%(i)s] = r_idx[%(i)s] * st[%(i)s];
                  r_end[%(i)s] = r_st[%(i)s] + ws[%(i)s];
                  // skip the padding
                  r_st[%(i)s] = r_st[%(i)s] < pd[%(i)s] ? pd[%(i)s] : r_st[%(i)s];
                  r_end[%(i)s] = r_end[%(i)s] > (r[%(i)s] - pd[%(i)s]) ? r[%(i)s] - pd[%(i)s] : r_end[%(i)s];
                  // from padded_img space to img space
                  r_st[%(i)s] -= pd[%(i)s];
                  r_end[%(i)s] -= pd[%(i)s];
                  // handle the case where no padding, ignore border is True
                  if (%(params)s->ignore_border)
                  {
                    r_end[%(i)s] = r_end[%(i)s] > r[%(i)s] ? r[%(i)s] : r_end[%(i)s];
                  }
                  // use the index to find the correct position in the output
                  o_idx[%(non_pool_ndim)s + %(i)s] = r_idx[%(i)s];
            """ % dict(
                i=i, non_pool_ndim=non_pool_ndim, params=sub["params"]
            )

        ccode += """
                  // get a pointer to the correct position in the output
                  dtype_%(z)s * z;
                  if (%(total_ndim)s == 4)
                    z = ((dtype_%(z)s*)(PyArray_GETPTR4(%(z)s, o_idx[0], o_idx[1], o_idx[2], o_idx[3])));
                  else
                    z = ((dtype_%(z)s*)(PyArray_GetPtr(%(z)s, o_idx)));
        """

        if self.mode == "max":
            for i in range(nd):
                ccode += """
                  // set the first index of dimension %(i)s
                  i_idx[%(non_pool_ndim)s + %(i)s] = r_st[%(i)s];
                """ % dict(
                    i=i, non_pool_ndim=non_pool_ndim
                )
            ccode += """
                  // use the first element as the initial value of collector
                  if (%(total_ndim)s == 4)
                    collector = ((dtype_%(x)s*)(PyArray_GETPTR4(%(x)s,i_idx[0],i_idx[1],i_idx[2],i_idx[3])))[0];
                  else
                    collector = ((dtype_%(x)s*)(PyArray_GetPtr(%(x)s,i_idx)))[0];
            """
            for i in range(nd):
                ccode += """
                  // go through the pooled region in the unpadded input
                  for(npy_intp m%(i)s=r_st[%(i)s]; m%(i)s<r_end[%(i)s]; m%(i)s++)
                  {
                    i_idx[%(non_pool_ndim)s + %(i)s] = m%(i)s;
                """ % dict(
                    i=i, non_pool_ndim=non_pool_ndim
                )
            ccode += """
                    // update maximum
                    dtype_%(x)s a;
                    if (%(total_ndim)s == 4)
                      a = ((dtype_%(x)s*)(PyArray_GETPTR4(%(x)s,i_idx[0],i_idx[1],i_idx[2],i_idx[3])))[0];
                    else
                      a = ((dtype_%(x)s*)(PyArray_GetPtr(%(x)s,i_idx)))[0];
                    collector = (a > collector) ? a : collector;
            """
            for i in range(nd):
                ccode += """
                  } // for loop over region
                """
            ccode += """
                  z[0] = collector;
            """
        elif self.mode in ("sum", "average_exc_pad", "average_inc_pad"):
            ccode += """
                  // initialize the sum at zero
                  collector = ((dtype_%(x)s)(0));
            """
            for i in range(nd):
                ccode += """
                  // go through the pooled region in the unpadded input
                  for(npy_intp m%(i)s=r_st[%(i)s]; m%(i)s<r_end[%(i)s]; m%(i)s++)
                  {
                    i_idx[%(non_pool_ndim)s + %(i)s] = m%(i)s;
                """ % dict(
                    i=i, non_pool_ndim=non_pool_ndim
                )
            ccode += """
                    // update sum
                    dtype_%(x)s a;
                    if (%(total_ndim)s == 4)
                      a = ((dtype_%(x)s*)(PyArray_GETPTR4(%(x)s,i_idx[0],i_idx[1],i_idx[2],i_idx[3])))[0];
                    else
                      a = ((dtype_%(x)s*)(PyArray_GetPtr(%(x)s,i_idx)))[0];
                    collector += a;
            """
            for i in range(nd):
                ccode += """
                  } // for loop over region
                """
            if self.mode == "sum":
                ccode += """
                  z[0] = collector;
                """
            elif self.mode == "average_inc_pad" and self.ignore_border:
                # region size = product over all pooling dimensions
                region_size = " * ".join(f"ws[{i}]" for i in range(nd))
                ccode += """
                  z[0] = collector / (%(region_size)s);
                """ % dict(
                    region_size=region_size
                )
            else:
                # region size = number elements of in this region
                region_size = " * ".join(f"(r_end[{i}]-r_st[{i}])" for i in range(nd))
                ccode += """
                  z[0] = collector / (%(region_size)s);
                """ % dict(
                    region_size=region_size
                )
        for i in range(nd):
            ccode += """
            } // loop over pooling dimension
            """

        ccode += """
          } // for loop over non-pooling dimensions
        } // if z_prod
        """
        return ccode % locals()

    def c_code_cache_version(self):
        return (10, self.openmp)


class PoolGrad(OpenMPOp):
    __props__ = ("ignore_border", "mode", "ndim")

    @staticmethod
    def out_shape(
        imgshape,
        ws=None,
        ignore_border=False,
        stride=None,
        pad=None,
        ndim=2,
        ds=None,
        st=None,
        padding=None,
    ):
        """Return the shape of the output from this op, for input of given
        shape and flags.

        Parameters
        ----------
        imgshape : tuple of integers or scalar Aesara variables
            the shape of a tensor of images. The last N elements are
            interpreted as the downsampling dimensions.
        ws : tuple of N ints
            downsample factor over rows and columns this parameter
            indicates the size of the pooling region
        ignore_border : bool
            If ws doesn't divide imgshape, do we include an extra row/col/slice
            of partial downsampling (False) or ignore it (True).
        stride : list or tuple of N ints or None
            Stride size, which is the number of shifts over rows/cols/slices to get the
            next pool region. If stride is None, it is considered equal to ws
            (no overlap on pooling regions).
        pad : tuple of N ints or None
            For each downsampling dimension, this specifies the number of zeros to
            add as padding on both sides. For 2D and (pad_h, pad_w), pad_h specifies the
            size of the top and bottom margins, pad_w specifies the size of the left and
            right margins. No padding is added if pad is None.
        ndim : int
            The number of pooling dimensions N.
            The default is 2.
        ds
            *deprecated*, use parameter ws instead.
        st
            *deprecated*, use parameter st instead.
        padding
            *deprecated*, use parameter pad instead.

        Returns
        -------
        list :
            the shape of the output from this op, for input of given
            shape.  This will have the same length as imgshape, but
            with last N elements reduced as per the downsampling &
            ignore_border flags.

        """
        # check for deprecated parameter names
        if ds is not None:
            if ws is not None:
                raise ValueError(
                    "You can't provide a tuple value to both 'ws' and 'ds'."
                    " Please provide a value only to 'ws'."
                )
            else:
                warnings.warn(
                    "The 'ds' parameter in PoolGrad is not going"
                    " to exist anymore as it is going to be replaced by the"
                    " parameter 'ws'.",
                    category=DeprecationWarning,
                    stacklevel=2,
                )
                ws = ds
        elif ds is None and ws is None:
            raise ValueError("You must provide a tuple value for the window size.")

        if st is not None:
            if stride is not None:
                raise ValueError(
                    "You can't provide a tuple value to both 'st and 'stride'."
                    " Please provide a value only to 'stride'."
                )
            else:
                warnings.warn(
                    "The 'st' parameter in PoolGrad is not going"
                    " to exist anymore as it is going to be replaced by the"
                    " parameter 'stride'.",
                    category=DeprecationWarning,
                    stacklevel=2,
                )
                stride = st

        if padding is not None:
            if pad is not None:
                raise ValueError(
                    "You can't provide a tuple value to both 'padding' and pad."
                    "  Please provide a value only to pad."
                )
            else:
                warnings.warn(
                    "The 'padding' parameter in PoolGrad is not"
                    " going to exist anymore as it is going to be replaced"
                    " by the parameter 'pad'.",
                    category=DeprecationWarning,
                    stacklevel=2,
                )
                pad = padding

        if len(imgshape) < ndim:
            raise TypeError(f"imgshape must have at least {ndim} dimensions")

        if stride is None:
            stride = ws
        if pad is None:
            pad = (0,) * ndim
        patch_shape = tuple(
            at.extract_constant(imgshape[-ndim + i]) + pad[i] * 2 for i in range(ndim)
        )

        def compute_out(v, downsample, stride):
            if ignore_border:
                out = (v - downsample) // stride + 1
                if isinstance(out, Variable):
                    return tm.maximum(out, 0)
                else:
                    return np.maximum(out, 0)
            else:
                if isinstance(v, Variable):
                    return at.switch(
                        tm.ge(stride, downsample),
                        (v - 1) // stride + 1,
                        tm.maximum(0, (v - 1 - downsample) // stride + 1) + 1,
                    )
                elif stride >= downsample:
                    return (v - 1) // stride + 1
                else:
                    return max(0, (v - 1 - downsample) // stride + 1) + 1

        out_shape = [compute_out(patch_shape[i], ws[i], stride[i]) for i in range(ndim)]

        rval = list(imgshape[:-ndim]) + out_shape
        return rval

    def __init__(self, ignore_border, mode="max", ndim=2, openmp=None):
        self.ndim = ndim
        self.ignore_border = ignore_border
        if mode == "max_deterministic":
            # It seems max pool grad algo is already deterministic in CPU.
            mode = "max"
        if mode not in ("max", "sum", "average_inc_pad", "average_exc_pad"):
            raise ValueError(
                "Pool mode parameter only support 'max', 'sum',"
                " 'average_inc_pad' and 'average_exc_pad'. Got {mode}"
            )
        self.mode = mode
        super().__init__(openmp=openmp)

    def prepare_node(self, node, storage_map, compute_map, impl):
        if len(node.inputs) < 5:  # 5 for AveragePoolGrad, 6 for MaxPoolGrad
            # Old interface
            self.ndim = len(node.op.ds)
            self.mode = node.op.mode
            ws = at.constant(node.op.ds)
            st = at.constant(node.op.st)
            pad = at.constant(node.op.padding)
            node.inputs.append(ws)
            node.inputs.append(st)
            node.inputs.append(pad)
            if isinstance(ws, Constant):
                storage_map[ws] = [ws.data]
                compute_map[ws] = [True]
            else:
                storage_map[ws] = [None]
                compute_map[ws] = [False]
            if isinstance(st, Constant):
                storage_map[st] = [st.data]
                compute_map[st] = [True]
            else:
                storage_map[st] = [None]
                compute_map[st] = [False]
            if isinstance(pad, Constant):
                storage_map[pad] = [pad.data]
                compute_map[pad] = [True]
            else:
                storage_map[pad] = [None]
                compute_map[pad] = [False]

    def infer_shape(self, fgraph, node, in_shapes):
        return [in_shapes[0]]


class MaxPoolGrad(PoolGrad):
    # params_type ignore_border don't change c code

    def __init__(self, ignore_border, ndim=2, openmp=None):
        PoolGrad.__init__(self, ignore_border, mode="max", ndim=ndim, openmp=openmp)

    def make_node(self, x, maxout, gz, ws, stride=None, pad=None):
        # make_node should only be called by the grad function of
        # Pool, so these asserts should not fail.
        x = at.as_tensor_variable(x)
        maxout = at.as_tensor_variable(maxout)
        gz = at.as_tensor_variable(gz)
        nd = self.ndim
        if stride is None:
            stride = ws
        if pad is None:
            pad = (0,) * nd
        ws = at.as_tensor_variable(ws)
        stride = at.as_tensor_variable(stride)
        pad = at.as_tensor_variable(pad)
        assert isinstance(x, Variable) and x.ndim >= nd
        assert isinstance(maxout, Variable) and maxout.ndim >= nd
        assert isinstance(gz, Variable) and gz.ndim >= nd
        assert isinstance(ws, Variable) and ws.ndim == 1
        assert isinstance(stride, Variable) and stride.ndim == 1
        assert isinstance(pad, Variable) and pad.ndim == 1
        assert x.ndim == maxout.ndim == gz.ndim >= nd
        if ws.dtype not in int_dtypes:
            raise TypeError("Pool downsample parameters must be ints.")
        if stride.dtype not in int_dtypes:
            raise TypeError("Stride parameters must be ints.")
        if pad.dtype not in int_dtypes:
            raise TypeError("Padding parameters must be ints.")
        return Apply(self, [x, maxout, gz, ws, stride, pad], [x.type()])

    def perform(self, node, inp, out):
        assert self.mode == "max"
        x, maxout, gz, ws, stride, pad = inp
        (gx_stg,) = out
        nd = self.ndim
        assert ws.shape == stride.shape == pad.shape == (nd,)
        if len(x.shape) < nd:
            raise NotImplementedError(
                f"MaxPoolGrad requires input with {nd} or more dimensions"
            )
        pool_out_shp = maxout.shape[-nd:]
        img_shp = tuple(x.shape[-nd + i] + 2 * pad[i] for i in range(nd))

        # pad the image
        if max(pad) != 0:
            y = np.zeros(x.shape[:-nd] + img_shp, dtype=x.dtype)
            y[
                (slice(None),) * (len(x.shape) - nd)
                + tuple(slice(pad[i], img_shp[i] - pad[i]) for i in range(nd))
            ] = x
        else:
            y = x
        gx = np.zeros_like(y)

        # precompute the region boundaries for each dimension
        region_ranges = [[] for i in range(nd)]
        for i in range(nd):
            for j in range(pool_out_shp[i]):
                start = max(j * stride[i], pad[i])
                end = min(start + ws[i], img_shp[i])
                region_ranges[i].append(range(start, end))

        # iterate over non-pooling dimensions
        for k in np.ndindex(*x.shape[:-nd]):
            gxk = gx[k]
            gzk = gz[k]
            yk = y[k]
            maxoutk = maxout[k]
            # iterate over pooling regions
            for r in np.ndindex(*pool_out_shp):
                maxout_value = maxoutk[r]
                # iterate inside region
                for c in itertools.product(
                    *[region_ranges[i][r[i]] for i in range(nd)]
                ):
                    if maxout_value == yk[c]:
                        gxk[c] += gzk[r]

        # unpad the image
        gx = gx[
            (slice(None),) * (len(x.shape) - nd)
            + tuple(slice(pad[i], img_shp[i] - pad[i]) for i in range(nd))
        ]
        gx_stg[0] = gx

    def grad(self, inp, grads):
        x, maxout, gz, ws, stride, pad = inp
        (ggx,) = grads
        return [
            at.zeros_like(x),
            at.zeros_like(maxout),
            DownsampleFactorMaxGradGrad(
                ndim=self.ndim, ignore_border=self.ignore_border
            )(x, maxout, ggx, ws, stride, pad),
        ] + [DisconnectedType()() for i in inp[3:]]

    def connection_pattern(self, node):
        return [[1], [1], [1], [0], [0], [0]]

    def c_code(self, node, name, inp, out, sub):
        assert self.mode == "max"
        x, z, gz, ws, stride, pad = inp
        (gx,) = out
        nd = self.ndim
        total_ndim = node.inputs[0].ndim
        non_pool_ndim = total_ndim - nd
        fail = sub["fail"]

        if self.openmp:
            # run in parallel over each pooling block
            omp_parallel = "#pragma omp parallel for private(r_st, r_end, r_idx, i_idx, o_idx, maximum) schedule(static)"
        else:
            omp_parallel = ""

        ccode = """
        // sanity checks
        int x_typenum = PyArray_ObjectType((PyObject*)%(x)s, 0);
        int z_typenum = PyArray_ObjectType((PyObject*)%(z)s, 0);
        int gz_typenum = PyArray_ObjectType((PyObject*)%(gz)s, 0);
        if ((x_typenum != z_typenum) || (x_typenum != gz_typenum))
        {
            PyErr_SetString(PyExc_ValueError, "input types must all match");
            %(fail)s;
        }
        if(PyArray_NDIM(%(x)s)!=%(total_ndim)s)
        {
            PyErr_SetString(PyExc_ValueError, "x must be a %(total_ndim)sD ndarray");
            %(fail)s;
        }
        if(PyArray_NDIM(%(z)s)!=%(total_ndim)s)
        {
            PyErr_SetString(PyExc_ValueError, "z must be a %(total_ndim)sD ndarray");
            %(fail)s;
        }
        if(PyArray_NDIM(%(gz)s)!=%(total_ndim)s)
        {
            PyErr_SetString(PyExc_ValueError, "gz must be a %(total_ndim)sD ndarray");
            %(fail)s;
        }
        if(PyArray_DIM(%(ws)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "ws must be a vector of size %(nd)s");
            %(fail)s;
        }
        if(PyArray_DIM(%(stride)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "stride must be a vector of size %(nd)s");
            %(fail)s;
        }
        if(PyArray_DIM(%(pad)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "pad must be a vector of size %(nd)s");
            %(fail)s;
        }
        npy_intp z[%(nd)s]; // shape of the output
        npy_intp r[%(nd)s]; // shape of the padded_input
        npy_intp ws[%(nd)s];
        npy_intp st[%(nd)s];
        npy_intp pd[%(nd)s];
        int nonzero_padding;
        nonzero_padding = 0;
        for (int i=0; i<%(nd)s; i++)
        {
            ws[i] = *((dtype_%(ws)s*)PyArray_GETPTR1(%(ws)s, i));
            st[i] = *((dtype_%(stride)s*)PyArray_GETPTR1(%(stride)s, i));
            pd[i] = *((dtype_%(pad)s*)PyArray_GETPTR1(%(pad)s, i));
            z[i] = PyArray_DIMS(%(z)s)[%(non_pool_ndim)s + i];
            r[i] = PyArray_DIMS(%(x)s)[%(non_pool_ndim)s + i] + 2 * pd[i];
            if (pd[i]>0)
                nonzero_padding = 1;
        }
        // allocating memory for output, if necessary
        int mem_nec;
        mem_nec = 0;
        if ((!%(gx)s) || !PyArray_ISCONTIGUOUS(%(gx)s)
            || *PyArray_DIMS(%(gx)s)!=%(total_ndim)s)
        {
            mem_nec = 1;
        }
        if (!mem_nec)
        {
            for (int i=0; i<%(total_ndim)s; i++)
            {
                if (PyArray_DIMS(%(gx)s)[i] != PyArray_DIMS(%(x)s)[i])
                {
                    mem_nec = 1;
                    break;
                }
            }
        }
        if (mem_nec)
        {
          Py_XDECREF(%(gx)s);
          %(gx)s = (PyArrayObject*) PyArray_ZEROS(%(total_ndim)s, PyArray_DIMS(%(x)s), x_typenum,0);
        }
        else {
          PyArray_FILLWBYTE(%(gx)s, 0);
        }
        dtype_%(z)s maximum; // temp var for maximum value in a region
        npy_intp z_prod;
        // do not run if any z[i] is zero
        z_prod = 1;
        for (int i=0; i<%(nd)s; i++)
        {
            z_prod *= z[i];
        }
        if (z_prod)
        {
            // will be used to hold start and end index of a region
            npy_intp r_st[%(nd)s];
            npy_intp r_end[%(nd)s];
            // index for iterating over the pooling regions
            npy_intp r_idx[%(nd)s];
            // placeholder for PyArray indexing (output)
            npy_intp o_idx[%(total_ndim)s];
            // placeholder for PyArray indexing (input)
            npy_intp i_idx[%(total_ndim)s];
            // loop over non-pooling dimensions
            npy_intp non_pooling_prod = 1;
            for (int i=0; i<%(non_pool_ndim)s; i++)
            {
                non_pooling_prod *= PyArray_DIMS(%(x)s)[i];
            }
            %(omp_parallel)s
            // first loop over non-pooling dimensions
            for (npy_intp t=0; t<non_pooling_prod; t++)
            {
                // compute the non-pooling index in each dimension
                if (%(non_pool_ndim)s!=0)
                {
                    o_idx[0] = t;
                    i_idx[0] = t;
                    for (int i=1; i<%(non_pool_ndim)s; i++)
                    {
                        o_idx[i] = o_idx[i - 1] / PyArray_DIMS(%(x)s)[i - 1];
                        o_idx[i - 1] =o_idx[i - 1] %% PyArray_DIMS(%(x)s)[i - 1];
                        i_idx[i] = o_idx[i];
                        i_idx[i - 1] = o_idx[i - 1];
                    }
                }

                // then loop over each region in each pooling dimension
        """

        for i in range(nd):
            ccode += """
                for (r_idx[%(i)s]=0; r_idx[%(i)s] < z[%(i)s]; r_idx[%(i)s]++) {
                  r_st[%(i)s] = r_idx[%(i)s] * st[%(i)s];
                  r_end[%(i)s] = r_st[%(i)s] + ws[%(i)s];
                  // skip the padding
                  r_st[%(i)s] = r_st[%(i)s] < pd[%(i)s] ? pd[%(i)s] : r_st[%(i)s];
                  r_end[%(i)s] = r_end[%(i)s] > (r[%(i)s] - pd[%(i)s]) ? r[%(i)s] - pd[%(i)s] : r_end[%(i)s];
                  // from padded_img space to img space
                  r_st[%(i)s] -= pd[%(i)s];
                  r_end[%(i)s] -= pd[%(i)s];
                  // use the index to find the correct position in the output
                  o_idx[%(non_pool_ndim)s + %(i)s] = r_idx[%(i)s];
            """ % dict(
                i=i, non_pool_ndim=non_pool_ndim
            )

        ccode += """
                  dtype_%(gz)s * gz;
                  if (%(total_ndim)s == 4)
                  {
                    // the maximum value
                    maximum = ((dtype_%(z)s*)(PyArray_GETPTR4(%(z)s,o_idx[0],o_idx[1],o_idx[2],o_idx[3])))[0];
                    // the gradient corresponding to this maximum value in z
                    gz = ((dtype_%(gz)s*)(PyArray_GETPTR4(%(gz)s, o_idx[0],o_idx[1],o_idx[2],o_idx[3])));
                  }
                  else
                  {
                    // the maximum value
                    maximum = ((dtype_%(z)s*)(PyArray_GetPtr(%(z)s,o_idx)))[0];
                    // the gradient corresponding to this maximum value in z
                    gz = ((dtype_%(gz)s*)(PyArray_GetPtr(%(gz)s, o_idx)));
                  }
        """
        for i in range(nd):
            ccode += """
                  // go through the pooled region in the unpadded input
                  for(npy_intp m%(i)s=r_st[%(i)s]; m%(i)s<r_end[%(i)s]; m%(i)s++)
                  {
                    i_idx[%(non_pool_ndim)s + %(i)s] = m%(i)s;
                """ % dict(
                i=i, non_pool_ndim=non_pool_ndim
            )
        ccode += """
                    dtype_%(x)s a;
                    dtype_%(gx)s * gx;
                    if (%(total_ndim)s == 4)
                    {
                      a = ((dtype_%(x)s*)(PyArray_GETPTR4(%(x)s,i_idx[0],i_idx[1],i_idx[2],i_idx[3])))[0];
                      gx = ((dtype_%(gx)s*)(PyArray_GETPTR4(%(gx)s, i_idx[0],i_idx[1],i_idx[2],i_idx[3])));
                    }
                    else
                    {
                      a = ((dtype_%(x)s*)(PyArray_GetPtr(%(x)s,i_idx)))[0];
                      gx = ((dtype_%(gx)s*)(PyArray_GetPtr(%(gx)s, i_idx)));
                    }
                    if (a == maximum){
                      gx[0] = gx[0] + gz[0];
                    }
        """
        for i in range(nd):
            ccode += """
                  } // for loop over region
                """
        for i in range(nd):
            ccode += """
                } // loop over pooling dimension
            """

        ccode += """
            } // for loop over non-pooling dimensions
        } // if z_prod
        """
        return ccode % locals()

    def c_code_cache_version(self):
        return (0, 11, self.openmp)


class AveragePoolGrad(PoolGrad):
    # ignore_border is used for perform, but not c code. No need in params_type

    def __init__(self, ignore_border, mode="average_inc_pad", ndim=2):
        assert mode in ("sum", "average_inc_pad", "average_exc_pad")
        PoolGrad.__init__(self, ignore_border, mode, ndim)

    # There is an extra dummy parameter to match the parameter count
    # of MaxPoolGrad.  They have to keep the same interface because of
    # the DownsampleFactorMaxGrad trick to keep old scripts working
    # (see downsample.py for details on this).
    def make_node(self, x, gz, ws, stride=None, pad=None, dummy=None):
        # make_node should only be called by the grad function of
        # Pool, so these asserts should not fail.
        x = at.as_tensor_variable(x)
        gz = at.as_tensor_variable(gz)
        nd = self.ndim
        if stride is None:
            stride = ws
        if pad is None:
            pad = (0,) * nd
        ws = at.as_tensor_variable(ws)
        stride = at.as_tensor_variable(stride)
        pad = at.as_tensor_variable(pad)
        assert isinstance(x, Variable) and x.ndim >= nd
        assert isinstance(gz, Variable) and gz.ndim >= nd
        assert isinstance(ws, Variable) and ws.ndim == 1
        assert isinstance(stride, Variable) and stride.ndim == 1
        assert x.ndim == gz.ndim >= nd
        assert isinstance(pad, Variable) and pad.ndim == 1
        if ws.dtype not in int_dtypes:
            raise TypeError("Pool downsample parameters must be ints.")
        if stride.dtype not in int_dtypes:
            raise TypeError("Stride parameters must be ints.")
        if pad.dtype not in int_dtypes:
            raise TypeError("Padding parameters must be ints.")
        return Apply(self, [x, gz, ws, stride, pad], [x.type()])

    def perform(self, node, inp, out):
        x, gz, ws, stride, pad = inp
        (gx_stg,) = out
        nd = self.ndim
        assert ws.shape == stride.shape == pad.shape == (nd,)
        if len(x.shape) < nd:
            raise NotImplementedError(
                f"AveragePoolGrad requires input with {nd} or more dimensions"
            )
        if self.mode == "average_exc_pad" and max(pad) != 0:
            raise NotImplementedError()
        z_shape = self.out_shape(x.shape, ws, self.ignore_border, stride, pad, nd)
        if (gx_stg[0] is None) or (gx_stg[0].shape != z_shape):
            gx_stg[0] = np.empty(z_shape, dtype=x.dtype)
        zz = gx_stg[0]
        # size of pooling output
        pool_out_shp = zz.shape[-nd:]
        img_shp = tuple(x.shape[-nd + i] + 2 * pad[i] for i in range(nd))
        inc_pad = self.mode == "average_inc_pad"
        sum_mode = self.mode == "sum"

        # initialize the padded output
        gx = np.zeros((x.shape[:-nd] + img_shp), dtype=x.dtype)

        # precompute the region boundaries and sizes for each dimension
        region_slices = [[] for i in range(nd)]
        region_sizes = [[] for i in range(nd)]
        for i in range(nd):
            for j in range(pool_out_shp[i]):
                if sum_mode or inc_pad:
                    start = j * stride[i]
                else:
                    start = max(j * stride[i], pad[i])
                end = min(start + ws[i], img_shp[i])
                region_slices[i].append(slice(start, end))
                region_sizes[i].append(end - start)

        # iterate over non-pooling dimensions
        region_slice = [None] * nd
        for k in np.ndindex(*x.shape[:-nd]):
            gzk = gz[k]
            gxk = gx[k]
            # iterate over pooling regions
            for r in np.ndindex(*pool_out_shp):
                region_size = 1
                for i in range(nd):
                    region_slice[i] = region_slices[i][r[i]]
                    region_size *= region_sizes[i][r[i]]
                if sum_mode:
                    val = gzk[r]
                else:
                    # divide by region size
                    val = gzk[r] / region_size
                gxk[region_slice] += val

        # unpad the image
        gx = gx[
            (slice(None),) * (len(x.shape) - nd)
            + tuple(slice(pad[i], img_shp[i] - pad[i]) for i in range(nd))
        ]
        gx_stg[0] = gx

    def grad(self, inp, grads):
        x, gz, ws, stride, pad = inp
        (ggx,) = grads
        return [
            at.zeros_like(x),
            Pool(ignore_border=self.ignore_border, ndim=self.ndim, mode=self.mode)(
                ggx, ws, stride, pad
            ),
        ] + [DisconnectedType()() for i in inp[2:]]

    def connection_pattern(self, node):
        return [[1], [1], [0], [0], [0]]

    def c_code(self, node, name, inp, out, sub):
        x, gz, ws, stride, pad = inp
        (gx,) = out
        nd = self.ndim
        total_ndim = node.inputs[0].ndim
        non_pool_ndim = total_ndim - nd
        fail = sub["fail"]
        inc_pad = int(self.mode == "average_inc_pad")
        sum_mode = int(self.mode == "sum")
        if self.openmp:
            # run in parallel over each pooling block
            omp_parallel = "#pragma omp parallel for private(r_st, r_end, r_pad_width, r_idx, i_idx, o_idx) schedule(static)"
        else:
            omp_parallel = ""

        ccode = """
        // sanity checks
        int x_typenum = PyArray_ObjectType((PyObject*)%(x)s, 0);
        int gz_typenum = PyArray_ObjectType((PyObject*)%(gz)s, 0);
        if (x_typenum != gz_typenum)
        {
            PyErr_SetString(PyExc_ValueError, "input types must all match");
            %(fail)s;
        }
        if(PyArray_NDIM(%(x)s)!=%(total_ndim)s)
        {
            PyErr_SetString(PyExc_ValueError, "x must be a %(total_ndim)sD ndarray");
            %(fail)s;
        }
        if(PyArray_NDIM(%(gz)s)!=%(total_ndim)s)
        {
            PyErr_SetString(PyExc_ValueError, "gz must be a %(total_ndim)sD ndarray");
            %(fail)s;
        }
        if(PyArray_DIM(%(ws)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "ws must be a vector of size %(nd)s");
            %(fail)s;
        }
        if(PyArray_DIM(%(stride)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "stride must be a vector of size %(nd)s");
            %(fail)s;
        }
        if(PyArray_DIM(%(pad)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "pad must be a vector of size %(nd)s");
            %(fail)s;
        }
        npy_intp z[%(nd)s]; // shape of the output
        npy_intp r[%(nd)s]; // shape of the padded_input
        npy_intp ws[%(nd)s];
        npy_intp st[%(nd)s];
        npy_intp pd[%(nd)s];
        int nonzero_padding;
        nonzero_padding = 0;
        for (int i=0; i<%(nd)s; i++)
        {
            ws[i] = *((dtype_%(ws)s*)PyArray_GETPTR1(%(ws)s, i));
            st[i] = *((dtype_%(stride)s*)PyArray_GETPTR1(%(stride)s, i));
            pd[i] = *((dtype_%(pad)s*)PyArray_GETPTR1(%(pad)s, i));
            z[i] = PyArray_DIMS(%(gz)s)[%(non_pool_ndim)s + i];
            r[i] = PyArray_DIMS(%(x)s)[%(non_pool_ndim)s + i] + 2 * pd[i];
            if (pd[i]>0)
                nonzero_padding = 1;
        }
        if (!%(inc_pad)s && !%(sum_mode)s && nonzero_padding)
        {
            PyErr_SetString(PyExc_ValueError,
              "padding must be zero for average_exc_pad");
            %(fail)s;
        }
        // allocating memory for output, if necessary
        int mem_nec;
        mem_nec = 0;
        if ((!%(gx)s) || !PyArray_ISCONTIGUOUS(%(gx)s)
            || *PyArray_DIMS(%(gx)s)!=%(total_ndim)s)
        {
            mem_nec = 1;
        }
        if (!mem_nec)
        {
            for (int i=0; i<%(total_ndim)s; i++)
            {
                if (PyArray_DIMS(%(gx)s)[i] != PyArray_DIMS(%(x)s)[i])
                {
                    mem_nec = 1;
                    break;
                }
            }
        }
        if (mem_nec)
        {
          Py_XDECREF(%(gx)s);
          %(gx)s = (PyArrayObject*) PyArray_ZEROS(%(total_ndim)s, PyArray_DIMS(%(x)s), x_typenum,0);
        }
        else {
          PyArray_FILLWBYTE(%(gx)s, 0);
        }
        npy_intp z_prod;
        // do not run if any z[i] is zero
        z_prod = 1;
        for (int i=0; i<%(nd)s; i++)
        {
            z_prod *= z[i];
        }
        if (z_prod)
        {
            // will be used to hold start and end index of a region
            npy_intp r_st[%(nd)s];
            npy_intp r_end[%(nd)s];
            // padded region size
            npy_intp r_pad_width[%(nd)s];
            // index for iterating over the pooling regions
            npy_intp r_idx[%(nd)s];
            // placeholder for PyArray indexing (output)
            npy_intp o_idx[%(total_ndim)s];
            // placeholder for PyArray indexing (input)
            npy_intp i_idx[%(total_ndim)s];
            // loop over non-pooling dimensions
            npy_intp non_pooling_prod = 1;
            for (int i=0; i<%(non_pool_ndim)s; i++)
            {
                non_pooling_prod *= PyArray_DIMS(%(x)s)[i];
            }
            %(omp_parallel)s
            // first loop over non-pooling dimensions
            for (npy_intp t=0; t<non_pooling_prod; t++)
            {
                // compute the non-pooling index in each dimension
                if (%(non_pool_ndim)s!=0)
                {
                    o_idx[0] = t;
                    i_idx[0] = t;
                    for (int i=1; i<%(non_pool_ndim)s; i++)
                    {
                        o_idx[i] = o_idx[i - 1] / PyArray_DIMS(%(x)s)[i - 1];
                        o_idx[i - 1] =o_idx[i - 1] %% PyArray_DIMS(%(x)s)[i - 1];
                        i_idx[i] = o_idx[i];
                        i_idx[i - 1] = o_idx[i - 1];
                    }
                }

                // then loop over each region in each pooling dimension
        """

        for i in range(nd):
            ccode += """
                for (r_idx[%(i)s]=0; r_idx[%(i)s] < z[%(i)s]; r_idx[%(i)s]++) {
                  r_st[%(i)s] = r_idx[%(i)s] * st[%(i)s];
                  if (!%(sum_mode)s && !%(inc_pad)s && r_st[%(i)s] < pd[%(i)s])
                  {
                    r_st[%(i)s] = pd[%(i)s];
                  }
                  r_end[%(i)s] = r_st[%(i)s] + ws[%(i)s];
                  r_end[%(i)s] = r_end[%(i)s] > r[%(i)s] ? r[%(i)s] : r_end[%(i)s];
                  r_pad_width[%(i)s] = r_end[%(i)s] - r_st[%(i)s];
                  // from padded_img space to img space
                  r_st[%(i)s] = r_st[%(i)s] - pd[%(i)s] > 0 ? r_st[%(i)s] - pd[%(i)s] : 0;
                  r_end[%(i)s] = r_end[%(i)s] > r[%(i)s] - pd[%(i)s] ? r[%(i)s] - 2 * pd[%(i)s] : r_end[%(i)s] - pd[%(i)s];

                  // use the index to find the correct position in the output
                  o_idx[%(non_pool_ndim)s + %(i)s] = r_idx[%(i)s];
            """ % dict(
                i=i, sum_mode=sum_mode, inc_pad=inc_pad, non_pool_ndim=non_pool_ndim
            )

        ccode += """
                  dtype_%(gz)s * gz;
                  dtype_%(gz)s val;
                  if (%(total_ndim)s == 4)
                  {
                    // the gradient for this region
                    gz = ((dtype_%(gz)s*)(PyArray_GETPTR4(%(gz)s, o_idx[0],o_idx[1],o_idx[2],o_idx[3])));
                  }
                  else
                  {
                    // the gradient for this region
                    gz = ((dtype_%(gz)s*)(PyArray_GetPtr(%(gz)s, o_idx)));
                  }
                  // compute the contribution
                  if (%(sum_mode)s)
                  {
                    val = gz[0];
                  }
                  else
                  {
                    val = gz[0] / (%(region_size)s);
                  }
        """
        region_size = " * ".join(f"r_pad_width[{i}]" for i in range(nd))
        for i in range(nd):
            ccode += """
                  // go through the pooled region in the unpadded input
                  for(npy_intp m%(i)s=r_st[%(i)s]; m%(i)s<r_end[%(i)s]; m%(i)s++)
                  {
                    i_idx[%(non_pool_ndim)s + %(i)s] = m%(i)s;
                """ % dict(
                i=i, non_pool_ndim=non_pool_ndim
            )
        ccode += """
                    dtype_%(gx)s * gx;
                    if (%(total_ndim)s == 4)
                    {
                      gx = ((dtype_%(gx)s*)(PyArray_GETPTR4(%(gx)s, i_idx[0],i_idx[1],i_idx[2],i_idx[3])));
                    }
                    else
                    {
                      gx = ((dtype_%(gx)s*)(PyArray_GetPtr(%(gx)s, i_idx)));
                    }
                    gx[0] = gx[0] + val;
        """
        for i in range(nd):
            ccode += """
                  } // for loop over region
                """
        for i in range(nd):
            ccode += """
                } // loop over pooling dimension
            """

        ccode += """
            } // for loop over non-pooling dimensions
        } // if z_prod
        """
        return ccode % locals()

    def c_code_cache_version(self):
        return (0, 4, self.openmp)


class DownsampleFactorMaxGradGrad(OpenMPOp):
    __props__ = ("ignore_border", "mode", "ndim")

    def __init__(self, ignore_border, mode="max", ndim=2, openmp=None):
        self.ndim = ndim
        self.ignore_border = ignore_border
        self.mode = mode
        super().__init__(openmp=openmp)
        assert self.mode == "max"

    def make_node(self, x, maxout, gz, ws, stride=None, pad=None):
        # make_node should only be called by the grad function of
        # MaxPoolGrad, so these asserts should not fail.
        x = at.as_tensor_variable(x)
        maxout = at.as_tensor_variable(maxout)
        gz = at.as_tensor_variable(gz)
        nd = self.ndim
        if stride is None:
            stride = ws
        if pad is None:
            pad = (0,) * nd
        elif isinstance(pad, (tuple, list)):
            if max(pad) != 0 and not self.ignore_border:
                raise NotImplementedError("padding works only with ignore_border=True")
            if isinstance(ws, (tuple, list)):
                if any(pad[i] >= ws[i] for i in range(nd)):
                    raise NotImplementedError("padding must be smaller than strides")
        ws = at.as_tensor_variable(ws)
        stride = at.as_tensor_variable(stride)
        pad = at.as_tensor_variable(pad)
        assert ws.ndim == 1
        assert stride.ndim == 1
        assert pad.ndim == 1
        assert x.ndim == maxout.ndim == gz.ndim >= nd
        if ws.dtype not in int_dtypes:
            raise TypeError("Pool downsample parameters must be ints.")
        if stride.dtype not in int_dtypes:
            raise TypeError("Stride parameters must be ints.")
        if pad.dtype not in int_dtypes:
            raise TypeError("Padding parameters must be ints.")
        return Apply(self, [x, maxout, gz, ws, stride, pad], [x.type()])

    def perform(self, node, inp, out):
        x, maxout, ggx, ws, stride, pad = inp
        (z,) = out
        nd = self.ndim
        assert ws.shape == stride.shape == pad.shape == (nd,)
        if len(x.shape) < nd:
            raise NotImplementedError(
                "DownsampleFactorMaxGradGrad requires input "
                "with {} or more dimensions".format(nd)
            )
        if (z[0] is None) or (z[0].shape != maxout.shape):
            z[0] = np.zeros(maxout.shape, dtype=x.dtype)
        ggz = z[0]  # grad wrt maxout_grad has the same shape as maxout
        # size of pooling output
        pool_out_shp = ggz.shape[-nd:]
        img_shp = tuple(x.shape[-nd + i] + 2 * pad[i] for i in range(nd))

        # pad the image and its gradients
        if max(pad) > 0:
            y_padded = np.zeros(x.shape[:-nd] + img_shp, dtype=x.dtype)
            y_padded[
                (slice(None),) * (len(x.shape) - nd)
                + tuple(slice(pad[i], img_shp[i] - pad[i]) for i in range(nd))
            ] = x
            ggx_padded = np.zeros(x.shape[:-nd] + img_shp, dtype=x.dtype)
            ggx_padded[
                (slice(None),) * (len(x.shape) - nd)
                + tuple(slice(pad[i], img_shp[i] - pad[i]) for i in range(nd))
            ] = ggx

        else:
            y_padded = x
            ggx_padded = ggx

        # precompute the region boundaries for each dimension
        region_ranges = [[] for i in range(nd)]
        for i in range(nd):
            for j in range(pool_out_shp[i]):
                start = j * stride[i]
                end = min(start + ws[i], img_shp[i])
                region_ranges[i].append(range(start, end))

        # iterate over non-pooling dimensions
        for k in np.ndindex(*x.shape[:-nd]):
            ggxk = ggx_padded[k]
            ggzk = ggz[k]
            yk = y_padded[k]
            maxoutk = maxout[k]
            # iterate over pooling regions
            for r in np.ndindex(*pool_out_shp):
                # iterate inside region
                maxout_value = maxoutk[r]
                for c in itertools.product(
                    *[region_ranges[i][r[i]] for i in range(nd)]
                ):
                    if maxout_value == yk[c]:
                        ggzk[r] += ggxk[c]

    def infer_shape(self, fgraph, node, in_shapes):
        return [in_shapes[1]]

    def grad(self, inp, grads):
        x, maxout, ggx, ws, stride, pad = inp
        (gz,) = grads
        return [
            at.zeros_like(x),
            at.zeros_like(maxout),
            MaxPoolGrad(ignore_border=self.ignore_border, ndim=self.ndim)(
                x, maxout, gz, ws, stride, pad
            ),
            DisconnectedType()(),
            DisconnectedType()(),
            DisconnectedType()(),
        ]

    def connection_pattern(self, node):
        return [[1], [1], [1], [0], [0], [0]]

    def c_code(self, node, name, inp, out, sub):
        if self.mode != "max":
            raise MethodNotDefined()
        x, maxout, ggx, ws, stride, pad = inp
        (z,) = out  # the grad of grad
        nd = self.ndim
        total_ndim = node.inputs[0].ndim
        non_pool_ndim = total_ndim - nd
        fail = sub["fail"]

        if self.openmp:
            # run in parallel over each pooling block
            omp_parallel = "#pragma omp parallel for private(r_st, r_end, r_idx, i_idx, o_idx, maximum) schedule(static)"
        else:
            omp_parallel = ""
        ccode = """
        int z_typenum = PyArray_ObjectType((PyObject*)%(maxout)s, 0);
        npy_intp z[%(nd)s]; // shape of the output
        npy_intp r[%(nd)s]; // shape of the padded_input
        npy_intp ws[%(nd)s];
        npy_intp st[%(nd)s];
        npy_intp pd[%(nd)s];
        if(PyArray_DIM(%(ws)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "ws must be a vector of size %(nd)s");
            %(fail)s;
        }
        if(PyArray_DIM(%(stride)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "stride must be a vector of size %(nd)s");
            %(fail)s;
        }
        if(PyArray_DIM(%(pad)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "pad must be a vector of size %(nd)s");
            %(fail)s;
        }
        for (int i=0; i<%(nd)s; i++)
        {
            ws[i] = *((dtype_%(ws)s*)PyArray_GETPTR1(%(ws)s, i));
            st[i] = *((dtype_%(stride)s*)PyArray_GETPTR1(%(stride)s, i));
            pd[i] = *((dtype_%(pad)s*)PyArray_GETPTR1(%(pad)s, i));
            z[i] = PyArray_DIMS(%(maxout)s)[%(non_pool_ndim)s + i];
            r[i] = PyArray_DIMS(%(x)s)[%(non_pool_ndim)s + i] + 2 * pd[i];
        }
        // allocating memory for output, if necessary
        int mem_nec;
        mem_nec = 0;
        if ((!%(z)s) || !PyArray_ISCONTIGUOUS(%(z)s)
            || *PyArray_DIMS(%(z)s)!=%(total_ndim)s)
        {
            mem_nec = 1;
        }
        if (!mem_nec)
        {
            for (int i=0; i<%(total_ndim)s; i++)
            {
                if (PyArray_DIMS(%(z)s)[i] != PyArray_DIMS(%(maxout)s)[i])
                {
                    mem_nec = 1;
                    break;
                }
            }
        }
        if (mem_nec)
        {
          Py_XDECREF(%(z)s);
          %(z)s = (PyArrayObject*) PyArray_ZEROS(%(total_ndim)s, PyArray_DIMS(%(maxout)s), z_typenum,0);
        }
        else {
          PyArray_FILLWBYTE(%(z)s, 0);
        }
        dtype_%(maxout)s maximum; // temp var for maximum value in a region
        // will be used to hold start and end index of a region
        npy_intp r_st[%(nd)s];
        npy_intp r_end[%(nd)s];
        // index for iterating over the pooling regions
        npy_intp r_idx[%(nd)s];
        // placeholder for PyArray indexing (output)
        npy_intp o_idx[%(total_ndim)s];
        // placeholder for PyArray indexing (input)
        npy_intp i_idx[%(total_ndim)s];
        // loop over non-pooling dimensions
        npy_intp non_pooling_prod;
        non_pooling_prod = 1;
        for (int i=0; i<%(non_pool_ndim)s; i++)
        {
            non_pooling_prod *= PyArray_DIMS(%(x)s)[i];
        }
        %(omp_parallel)s
        // first loop over non-pooling dimensions
        for (npy_intp t=0; t<non_pooling_prod; t++)
        {
            // compute the non-pooling index in each dimension
            if (%(non_pool_ndim)s!=0)
            {
                o_idx[0] = t;
                i_idx[0] = t;
                for (int i=1; i<%(non_pool_ndim)s; i++)
                {
                    o_idx[i] = o_idx[i - 1] / PyArray_DIMS(%(x)s)[i - 1];
                    o_idx[i - 1] = o_idx[i - 1] %% PyArray_DIMS(%(x)s)[i - 1];
                    i_idx[i] = o_idx[i];
                    i_idx[i - 1] = o_idx[i - 1];
                }
            }

            // then loop over each region in each pooling dimension
        """

        for i in range(nd):
            ccode += """
                for (r_idx[%(i)s]=0; r_idx[%(i)s] < z[%(i)s]; r_idx[%(i)s]++) {
                  r_st[%(i)s] = r_idx[%(i)s] * st[%(i)s];
                  r_end[%(i)s] = r_st[%(i)s] + ws[%(i)s];
                  // skip the padding
                  r_st[%(i)s] = r_st[%(i)s] < pd[%(i)s] ? pd[%(i)s] : r_st[%(i)s];
                  r_end[%(i)s] = r_end[%(i)s] > (r[%(i)s] - pd[%(i)s]) ? r[%(i)s] - pd[%(i)s] : r_end[%(i)s];
                  // from padded_img space to img space
                  r_st[%(i)s] -= pd[%(i)s];
                  r_end[%(i)s] -= pd[%(i)s];
                  // use the index to find the correct position in the output
                  o_idx[%(non_pool_ndim)s + %(i)s] = r_idx[%(i)s];
            """ % dict(
                i=i, non_pool_ndim=non_pool_ndim
            )

        ccode += """
                  dtype_%(z)s * z;
                  if (%(total_ndim)s == 4)
                  {
                    // the maximum value
                    maximum = ((dtype_%(maxout)s*)(PyArray_GETPTR4(%(maxout)s,o_idx[0],o_idx[1],o_idx[2],o_idx[3])))[0];
                    // z at this position
                    z = ((dtype_%(z)s*)(PyArray_GETPTR4(%(z)s,o_idx[0],o_idx[1],o_idx[2],o_idx[3])));
                  }
                  else
                  {
                    // the maximum value
                    maximum = ((dtype_%(maxout)s*)(PyArray_GetPtr(%(maxout)s,o_idx)))[0];
                    // z at this position
                    z = ((dtype_%(z)s*)(PyArray_GetPtr(%(z)s,o_idx)));
                  }
        """
        for i in range(nd):
            ccode += """
                  // go through the pooled region in the unpadded input
                  for(npy_intp m%(i)s=r_st[%(i)s]; m%(i)s<r_end[%(i)s]; m%(i)s++)
                  {
                    i_idx[%(non_pool_ndim)s + %(i)s] = m%(i)s;
                """ % dict(
                i=i, non_pool_ndim=non_pool_ndim
            )
        ccode += """
                    dtype_%(x)s a;
                    dtype_%(ggx)s * ggx;
                    if (%(total_ndim)s == 4)
                    {
                      a = ((dtype_%(x)s*)(PyArray_GETPTR4(%(x)s,i_idx[0],i_idx[1],i_idx[2],i_idx[3])))[0];
                      ggx = ((dtype_%(ggx)s*)(PyArray_GETPTR4(%(ggx)s,i_idx[0],i_idx[1],i_idx[2],i_idx[3])));
                    }
                    else
                    {
                      a = ((dtype_%(x)s*)(PyArray_GetPtr(%(x)s,i_idx)))[0];
                      ggx = ((dtype_%(ggx)s*)(PyArray_GetPtr(%(ggx)s,i_idx)));
                    }
                    if (a == maximum){
                      z[0] += ggx[0];
                    }
        """
        for i in range(nd):
            ccode += """
                  } // for loop over region
                """
        for i in range(nd):
            ccode += """
              } // loop over pooling dimension
            """

        ccode += """
          } // for loop over non-pooling dimensions
        """
        return ccode % locals()

    def c_code_cache_version(self):
        return (0, 5, self.openmp)


class MaxPoolRop(OpenMPOp):
    """
    Implements the R-operator for the downsample operation.

    Parameters
    ----------
    ws : list or tuple of N ints
        Downsample factor over rows, columns etc.
        ws indicates the size of the pooling region.
    ignore_border : bool
        If ws doesn't divide imgshape, do we include an extra row/col/slice
        of partial downsampling (False) or ignore it (True).
    stride : list or tuple of N ints or None
        Stride size, which is the number of shifts over rows/cols/slices to get the
        next pool region. If stride is None, it is considered equal to ws
        (no overlap on pooling regions).
    pad : tuple of N ints or None
        For each downsampling dimension, this specifies the number of zeros to
        add as padding on both sides. For 2D and (pad_h, pad_w), pad_h specifies the
        size of the top and bottom margins, pad_w specifies the size of the left and
        right margins. No padding is added if pad is None.
    mode : {'max', 'sum', 'average_inc_pad', 'average_exc_pad'}
        ('average_inc_pad' excludes the padding from the count,
        'average_exc_pad' include it)
    ndim : int
        The number of pooling dimensions N.
        The default is 2.
    """

    __props__ = ("ignore_border", "mode", "ndim")
    params_type = ParamsType(
        ignore_border=bool_t,
    )

    def __init__(self, ignore_border=False, mode="max", ndim=2, openmp=None):
        super().__init__(openmp=openmp)
        self.ndim = ndim
        self.ignore_border = ignore_border
        self.mode = mode
        assert mode == "max"

    def make_node(self, x, eval_point, ws, stride=None, pad=None):
        # TODO: consider restricting the dtype?
        x = at.as_tensor_variable(x)
        eval_point = at.as_tensor_variable(eval_point)
        nd = self.ndim
        if stride is None:
            stride = ws
        if pad is None:
            pad = (0,) * nd
        elif isinstance(pad, (tuple, list)):
            if max(pad) != 0 and not self.ignore_border:
                raise NotImplementedError("padding works only with ignore_border=True")
            if isinstance(ws, (tuple, list)):
                if any(pad[i] >= ws[i] for i in range(nd)):
                    raise NotImplementedError("padding must be smaller than strides")
        ws = at.as_tensor_variable(ws)
        stride = at.as_tensor_variable(stride)
        pad = at.as_tensor_variable(pad)
        assert ws.ndim == 1
        assert stride.ndim == 1
        assert pad.ndim == 1
        if x.type.ndim < nd:
            raise TypeError()
        if not ws.dtype.startswith("int"):
            raise TypeError("Pool downsample parameters must be ints.")
        if not stride.dtype.startswith("int"):
            raise TypeError("Stride parameters must be ints.")
        if not pad.dtype.startswith("int"):
            raise TypeError("Padding parameters must be ints.")
        # If the input shape are broadcastable we can have 0 in the output shape
        broad = x.broadcastable[:-nd] + (False,) * nd
        out = TensorType(eval_point.dtype, broad)
        return Apply(self, [x, eval_point, ws, stride, pad], [out()])

    def perform(self, node, inp, out, params):
        x, ex, ws, stride, pad = inp
        (z,) = out
        nd = self.ndim
        assert ws.shape == stride.shape == pad.shape == (nd,)
        if len(x.shape) < nd:
            raise NotImplementedError(
                f"Pool requires input with {nd} or more dimensions"
            )
        z_shape = Pool.out_shape(x.shape, ws, params.ignore_border, stride, pad, nd)
        if not self.ignore_border:
            assert all(z > 0 for z in z_shape[-nd:])
        if (z[0] is None) or (z[0].shape != z_shape):
            z[0] = np.empty(z_shape, dtype=x.dtype)
        zz = z[0]
        # size of pooling output
        pool_out_shp = zz.shape[-nd:]
        img_shp = tuple(x.shape[-nd + i] + 2 * pad[i] for i in range(nd))
        inc_pad = self.mode == "average_inc_pad"

        # pad the image and the eval point
        if max(pad) != 0:
            y = np.zeros(x.shape[:-nd] + img_shp, dtype=x.dtype)
            y[
                (slice(None),) * (len(x.shape) - nd)
                + tuple(slice(pad[i], img_shp[i] - pad[i]) for i in range(nd))
            ] = x
            ey = np.zeros(ex.shape[:-nd] + img_shp, dtype=ex.dtype)
            ey[
                (slice(None),) * (len(ex.shape) - nd)
                + tuple(slice(pad[i], img_shp[i] - pad[i]) for i in range(nd))
            ] = ex
        else:
            y = x
            ey = ex

        # precompute the region boundaries for each dimension
        region_slices = [[] for i in range(nd)]
        for i in range(nd):
            for j in range(pool_out_shp[i]):
                start = j * stride[i]
                end = min(start + ws[i], img_shp[i])
                if not inc_pad:
                    start = max(start, pad[i])
                    end = min(end, img_shp[i] - pad[i])
                region_slices[i].append(slice(start, end))

        # iterate over non-pooling dimensions
        for k in np.ndindex(*x.shape[:-nd]):
            zzk = zz[k]
            yk = y[k]
            eyk = ey[k]
            # iterate over pooling regions
            for r in np.ndindex(*pool_out_shp):
                # current slice in padded input
                ykslice = yk[[region_slices[i][r[i]] for i in range(nd)]]
                # current slice in eval points
                eykslice = eyk[[region_slices[i][r[i]] for i in range(nd)]]
                # indices of maximum
                idx = np.unravel_index(np.argmax(ykslice), ykslice.shape)
                zzk[r] = eykslice[idx]

    def c_headers(self, **kwargs):
        headers = ["<algorithm>"]
        headers += super().c_headers(**kwargs)
        return headers

    def c_code(self, node, name, inp, out, sub):
        if self.mode != "max":
            raise MethodNotDefined()
        x, ex, ws, stride, pad = inp
        (z,) = out
        nd = self.ndim
        total_ndim = node.inputs[0].ndim
        non_pool_ndim = total_ndim - nd
        fail = sub["fail"]
        params = sub["params"]

        if self.openmp:
            # run in parallel over each pooling block
            omp_parallel = "#pragma omp parallel for private(r_st, r_end, r_idx, i_idx, o_idx, collector, eval_collector) schedule(static)"
        else:
            omp_parallel = ""
        ccode = """
        int typenum = PyArray_ObjectType((PyObject*)%(x)s, 0);
        if(PyArray_NDIM(%(x)s)!=%(total_ndim)s)
        {
            PyErr_SetString(PyExc_ValueError, "x must be a %(total_ndim)sD ndarray");
            %(fail)s;
        }
        if(PyArray_NDIM(%(ex)s)!=%(total_ndim)s)
        {
            PyErr_SetString(PyExc_ValueError, "eval_point must be a %(total_ndim)sD ndarray");
            %(fail)s;
        }
        if(PyArray_DIM(%(ws)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "ws must be a vector of size %(nd)s");
            %(fail)s;
        }
        if(PyArray_DIM(%(stride)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "stride must be a vector of size %(nd)s");
            %(fail)s;
        }
        if(PyArray_DIM(%(pad)s, 0)!=%(nd)s)
        {
            PyErr_SetString(PyExc_ValueError, "pad must be a vector of size %(nd)s");
            %(fail)s;
        }
        npy_intp z[%(nd)s]; // shape of the output
        npy_intp r[%(nd)s]; // shape of the padded_input
        npy_intp ws[%(nd)s];
        npy_intp st[%(nd)s];
        npy_intp pd[%(nd)s];
        int nonzero_padding;
        nonzero_padding = 0;
        for (int i=0; i<%(nd)s; i++)
        {
            ws[i] = *((dtype_%(ws)s*)PyArray_GETPTR1(%(ws)s, i));
            st[i] = *((dtype_%(stride)s*)PyArray_GETPTR1(%(stride)s, i));
            pd[i] = *((dtype_%(pad)s*)PyArray_GETPTR1(%(pad)s, i));
            r[i] = PyArray_DIMS(%(x)s)[%(non_pool_ndim)s + i] + 2 * pd[i];
            if (pd[i]>0)
                nonzero_padding = 1;
        }
        if (!%(params)s->ignore_border && nonzero_padding)
        {
            PyErr_SetString(PyExc_ValueError,
              "padding must be zero when ignore border is False");
            %(fail)s;
        }
        if (%(params)s->ignore_border)
        {
            for (int i=0; i<%(nd)s; i++)
            {
                // '/' in C is different from '/' in python
                if (r[i] - ws[i] < 0)
                {
                  z[i] = 0;
                }
                else
                {
                  z[i] = (r[i] - ws[i]) / st[i] + 1;
                }
            }
        }
        else
        {
            for (int i=0; i<%(nd)s; i++)
            {
                // decide how many rows/cols the output has
                if (st[i] >= ws[i])
                {
                    z[i] = (r[i] - 1) / st[i] + 1;
                }
                else
                {
                    z[i] = std::max((npy_intp)0, (r[i] - 1 - ws[i] + st[i]) / st[i]) + 1;
                }
                assert(z[i] > 0);
            }
        }
        // memory allocation of z if necessary
        int mem_nec;
        mem_nec = 0;
        if ((!%(z)s) || *PyArray_DIMS(%(z)s)!=%(total_ndim)s)
        {
            mem_nec = 1;
        }
        if (!mem_nec)
        {
            for (int i=0; i<%(non_pool_ndim)s; i++)
            {
                if (PyArray_DIMS(%(z)s)[i] != PyArray_DIMS(%(x)s)[i])
                {
                    mem_nec = 1;
                    break;
                }
            }
        }
        if (!mem_nec)
        {
            for (int i=0; i<%(nd)s; i++)
            {
                if (PyArray_DIMS(%(z)s)[%(non_pool_ndim)s + i] != z[i])
                {
                    mem_nec = 1;
                    break;
                }
            }
        }
        if (mem_nec)
        {
          if (%(z)s) Py_XDECREF(%(z)s);
          npy_intp dims[%(total_ndim)s];
          for (int i=0; i<%(non_pool_ndim)s; i++)
          {
              dims[i] = PyArray_DIMS(%(x)s)[i];
          }
          for (int i=0; i<%(nd)s; i++)
          {
              dims[%(non_pool_ndim)s + i] = z[i];
          }
          //TODO: zeros not necessary
          %(z)s = (PyArrayObject*) PyArray_ZEROS(%(total_ndim)s, dims, typenum,0);
        }
        // initialize temp var for the value in a region
        dtype_%(x)s collector;
        dtype_%(ex)s eval_collector;
        npy_intp z_prod;
        // do not run if any z[i] is zero
        z_prod = 1;
        for (int i=0; i<%(nd)s; i++)
        {
            z_prod *= z[i];
        }
        if (z_prod)
        {
            // will be used to hold start and end index of a region
            npy_intp r_st[%(nd)s];
            npy_intp r_end[%(nd)s];
            // index for iterating over the pooling regions
            npy_intp r_idx[%(nd)s];
            // placeholder for PyArray indexing (output)
            npy_intp o_idx[%(total_ndim)s];
            // placeholder for PyArray indexing (input)
            npy_intp i_idx[%(total_ndim)s];
            // loop over non-pooling dimensions
            npy_intp non_pooling_prod = 1;
            for (int i=0; i<%(non_pool_ndim)s; i++)
            {
                non_pooling_prod *= PyArray_DIMS(%(x)s)[i];
            }
            %(omp_parallel)s
            // first loop over non-pooling dimensions
            for (npy_intp t=0; t<non_pooling_prod; t++)
            {
                // compute the non-pooling index in each dimension
                if (%(non_pool_ndim)s!=0)
                {
                    o_idx[0] = t;
                    i_idx[0] = t;
                    for (int i=1; i<%(non_pool_ndim)s; i++)
                    {
                        o_idx[i] = o_idx[i - 1] / PyArray_DIMS(%(x)s)[i - 1];
                        o_idx[i - 1] = o_idx[i - 1] %% PyArray_DIMS(%(x)s)[i - 1];
                        i_idx[i] = o_idx[i];
                        i_idx[i - 1] = o_idx[i - 1];
                    }
                }

                // then loop over each region in each pooling dimension
        """

        for i in range(nd):
            ccode += """
                for (r_idx[%(i)s]=0; r_idx[%(i)s] < z[%(i)s]; r_idx[%(i)s]++) {
                  r_st[%(i)s] = r_idx[%(i)s] * st[%(i)s];
                  r_end[%(i)s] = r_st[%(i)s] + ws[%(i)s];
                  // skip the padding
                  r_st[%(i)s] = r_st[%(i)s] < pd[%(i)s] ? pd[%(i)s] : r_st[%(i)s];
                  r_end[%(i)s] = r_end[%(i)s] > (r[%(i)s] - pd[%(i)s]) ? r[%(i)s] - pd[%(i)s] : r_end[%(i)s];
                  // from padded_img space to img space
                  r_st[%(i)s] -= pd[%(i)s];
                  r_end[%(i)s] -= pd[%(i)s];
                  // handle the case where no padding, ignore border is True
                  if (%(params)s->ignore_border)
                  {
                    r_end[%(i)s] = r_end[%(i)s] > r[%(i)s] ? r[%(i)s] : r_end[%(i)s];
                  }
                  // use the index to find the correct position in the output
                  o_idx[%(non_pool_ndim)s + %(i)s] = r_idx[%(i)s];
            """ % dict(
                i=i, params=sub["params"], non_pool_ndim=non_pool_ndim
            )

        ccode += """
                  // get a pointer to the correct position in the output
                  dtype_%(z)s * z;
                  if (%(total_ndim)s == 4)
                    z = ((dtype_%(z)s*)(PyArray_GETPTR4(%(z)s, o_idx[0], o_idx[1], o_idx[2], o_idx[3])));
                  else
                    z = ((dtype_%(z)s*)(PyArray_GetPtr(%(z)s, o_idx)));
        """

        for i in range(nd):
            ccode += """
              // set the first index of dimension %(i)s
              i_idx[%(non_pool_ndim)s + %(i)s] = r_st[%(i)s];
            """ % dict(
                i=i, non_pool_ndim=non_pool_ndim
            )
        ccode += """
              // use the first element as the initial value of collector
              if (%(total_ndim)s == 4) {
                collector = ((dtype_%(x)s*)(PyArray_GETPTR4(%(x)s,i_idx[0],i_idx[1],i_idx[2],i_idx[3])))[0];
                eval_collector = ((dtype_%(ex)s*)(PyArray_GETPTR4(%(ex)s,i_idx[0],i_idx[1],i_idx[2],i_idx[3])))[0];
              } else {
                collector = ((dtype_%(x)s*)(PyArray_GetPtr(%(x)s,i_idx)))[0];
                eval_collector = ((dtype_%(ex)s*)(PyArray_GetPtr(%(ex)s,i_idx)))[0];
              }
        """
        for i in range(nd):
            ccode += """
              // go through the pooled region in the unpadded input
              for(npy_intp m%(i)s=r_st[%(i)s]; m%(i)s<r_end[%(i)s]; m%(i)s++)
              {
                i_idx[%(non_pool_ndim)s + %(i)s] = m%(i)s;
            """ % dict(
                i=i, non_pool_ndim=non_pool_ndim
            )
        ccode += """
                // update maximum
                dtype_%(x)s a;
                dtype_%(ex)s ea;
                if (%(total_ndim)s == 4) {
                  a = ((dtype_%(x)s*)(PyArray_GETPTR4(%(x)s,i_idx[0],i_idx[1],i_idx[2],i_idx[3])))[0];
                  ea = ((dtype_%(ex)s*)(PyArray_GETPTR4(%(ex)s,i_idx[0],i_idx[1],i_idx[2],i_idx[3])))[0];
                }
                else {
                  a = ((dtype_%(x)s*)(PyArray_GetPtr(%(x)s,i_idx)))[0];
                  ea = ((dtype_%(ex)s*)(PyArray_GetPtr(%(ex)s,i_idx)))[0];
                }
                if (a > collector) {
                  collector = a;
                  eval_collector = ea;
                }
        """
        for i in range(nd):
            ccode += """
              } // for loop over region
            """
        ccode += """
              z[0] = eval_collector;
        """
        for i in range(nd):
            ccode += """
            } // loop over pooling dimension
            """

        ccode += """
          } // for loop over non-pooling dimensions
        } // if z_prod
        """
        return ccode % locals()

    def c_code_cache_version(self):
        return (2, self.openmp)
