import numpy

import chainer
from chainer import cuda
from chainer import function_node
from chainer import utils
from chainer.utils import type_check


if cuda.cudnn_enabled:
    cudnn = cuda.cudnn
    _mode = cudnn.cudnn.CUDNN_ACTIVATION_RELU


class ReLU(function_node.FunctionNode):

    """Rectified Linear Unit."""
    # TODO(beam2d): Implement in-place version.

    _use_cudnn = False

    def check_type_forward(self, in_types):
        type_check.expect(
            in_types.size() == 1,
            in_types[0].dtype.kind == 'f',
        )

    def forward_cpu(self, x):
        self.retain_outputs((0,))
        return utils.force_array(numpy.maximum(x[0], 0, dtype=x[0].dtype)),

    def forward_gpu(self, x):
        if chainer.should_use_cudnn('==always') and x[0].flags.c_contiguous:
            # cupy.activation_backward requires the input.
            # So, we retain it for backward computation.
            self.retain_inputs((0,))
            self._use_cudnn = True
            y = cudnn.activation_forward(x[0], _mode)
        else:
            y = cuda.cupy.maximum(x[0], 0)
        self.retain_outputs((0,))
        return y,

    def backward(self, indexes, gy):
        y = self.get_retained_outputs()[0]
        if chainer.should_use_cudnn('==always') and self._use_cudnn:
            x = self.get_retained_inputs()[0]
            return ReLUGrad3().apply((x, y, gy[0]))
        else:
            return ReLUGrad2().apply((y, gy[0]))


class Zero(function_node.FunctionNode):

    def forward(self, inputs):
        xp = chainer.cuda.get_array_module(*inputs)
        return utils.force_array(xp.zeros_like(inputs[0])),

    def backward(self, indexes, gy):
        return Zero().apply(gy)


class Heaviside(function_node.FunctionNode):

    def forward(self, inputs):
        x, = inputs
        return utils.force_array((x > 0).astype(x.dtype)),

    def backward(self, indexes, gy):
        return Zero().apply(gy)


class ReLUGrad2(function_node.FunctionNode):

    def forward_cpu(self, inputs):
        self.retain_inputs((0,))
        b, c = inputs
        y = (b > 0) * c
        return utils.force_array(y, dtype=y.dtype),

    def forward_gpu(self, inputs):
        self.retain_inputs((0,))
        b, c = inputs
        gx = cuda.elementwise(
            'T y, T gy', 'T gx',
            'gx = y > 0 ? gy : (T)0',
            'relu_bwd')(b, c)
        return gx,

    def backward(self, indexes, gy):
        ret = []
        if 0 in indexes:
            gb = Zero().apply(gy)[0]
            ret.append(gb)
        if 1 in indexes:
            b = self.get_retained_inputs()[0]
            gc = gy[0] * Heaviside().apply((b,))[0]
            ret.append(gc)
        return ret


class ReLUGrad3(function_node.FunctionNode):

    def forward_cpu(self, inputs):
        self.retain_inputs((0,))
        _, b, c = inputs
        y = (b > 0) * c
        return y,

    def forward_gpu(self, inputs):
        self.retain_inputs((1,))
        a, b, c = inputs
        assert chainer.should_use_cudnn('==always')
        y = cudnn.activation_backward(a, b, c, _mode)
        return y,

    def backward(self, indexes, gy):
        ret = []
        if 0 in indexes:
            ret.append(None)
        if 1 in indexes:
            gb = Zero().apply(gy)[0]
            ret.append(gb)
        if 2 in indexes:
            b = self.get_retained_inputs()[0]
            gc = gy[0] * Heaviside().apply((b,))[0]
            ret.append(gc)
        return ret


def relu(x):
    """Rectified Linear Unit function.

    .. math:: f(x)=\\max(0, x).

    Args:
        x (:class:`~chainer.Variable` or :class:`numpy.ndarray` or \
        :class:`cupy.ndarray`):
            Input variable. A :math:`(s_1, s_2, ..., s_N)`-shaped float array.

    Returns:
        ~chainer.Variable: Output variable. A
        :math:`(s_1, s_2, ..., s_N)`-shaped float array.

    .. admonition:: Example

        >>> x = np.array([[-1, 0], [2, -3], [-2, 1]], 'f')
        >>> np.any(x < 0)
        True
        >>> y = F.relu(x)
        >>> np.any(y.data < 0)
        False
        >>> y.shape
        (3, 2)

    """
    y, = ReLU().apply((x,))
    return y

@@ -1462,488 +1462,84 @@
 )%0A%0A%0A
-class Zero(function_node.FunctionNode):%0A%0A    def forward(self, inputs):%0A        xp = chainer.cuda.get_array_module(*inputs)%0A        return utils.force_array(xp.zeros_like(inputs%5B0%5D)),%0A%0A    def backward(self, indexes, gy):%0A        return Zero().apply(gy)%0A%0A%0Aclass Heaviside(function_node.FunctionNode):%0A%0A    def forward(self, inputs):%0A        x, = inputs%0A        return utils.force_array((x %3E 0).astype(x.dtype)),%0A%0A    def backward(self, indexes, gy):%0A        return Zero().apply(gy
+def _heaviside(x):%0A    return utils.force_array((x.data %3E 0).astype(x.dtype)
 )%0A%0A%0A
@@ -2076,45 +2076,8 @@
 es:%0A
-            gb = Zero().apply(gy)%5B0%5D%0A
     
@@ -2087,34 +2087,36 @@
      ret.append(
-gb
+None
 )%0A        if 1 i
@@ -2189,33 +2189,34 @@
    gc = gy%5B0%5D * 
-H
+_h
 eaviside().apply
@@ -2208,32 +2208,18 @@
 aviside(
+b
 )
-.apply((b,))%5B0%5D
 %0A       
@@ -2796,45 +2796,8 @@
 es:%0A
-            gb = Zero().apply(gy)%5B0%5D%0A
     
@@ -2819,10 +2819,12 @@
 end(
-gb
+None
 )%0A  
@@ -2921,9 +2921,10 @@
 %5D * 
-H
+_h
 eavi
@@ -2932,24 +2932,10 @@
 ide(
+b
 )
-.apply((b,))%5B0%5D
 %0A   

Simplify backward for efficient backprop