import numpy
import six

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


def _check_indices(indices):
    if len(indices) == 0:
        return
    # TODO(unno): Check indices without cpu
    indices = cuda.to_cpu(indices)
    for i in indices:
        if 0 <= i < len(indices):
            continue
        raise ValueError('Out of bounds index: {}'.format(i))
    sort = numpy.sort(indices)
    for s, t in six.moves.zip(sort, sort[1:]):
        if s == t:
            raise ValueError('indices contains duplicate value: {}'.format(s))


def _inverse_indices(indices):
    xp = cuda.get_array_module(indices)
    r = xp.empty_like(indices)
    if xp is numpy:
        r[indices] = numpy.arange(len(indices))
    else:
        cuda.elementwise(
            'S ind', 'raw S r',
            'r[ind] = i',
            'inverse_indices'
        )(indices, r)
    return r


class Permutate(function_node.FunctionNode):

    """Permutate function."""

    def __init__(self, axis=0, inv=False):
        self.axis = axis
        self.inv = inv

    def check_type_forward(self, in_types):
        type_check.argname(in_types, ('x', 'indices'))
        x_type, ind_type = in_types
        if self.axis < 0:
            type_check.expect(x_type.ndim >= -self.axis)
        else:
            type_check.expect(x_type.ndim > self.axis)

        type_check.expect(
            ind_type.dtype.kind == 'i',
            ind_type.ndim == 1,
            x_type.shape[self.axis] == ind_type.shape[0],
        )

    def _permutate(self, x, indices, inv):
        if inv:
            indices = _inverse_indices(indices)

        return x[((slice(None),) * self.axis) + (indices,)]

    def forward(self, inputs):
        self.retain_inputs((1,))
        x, inds = inputs

        if chainer.is_debug():
            _check_indices(inds)

        return self._permutate(x, inds, self.inv),

    def backward(self, indexes, grad_outputs):
        inds = self.inputs[1]
        g, = grad_outputs
        gx, = Permutate(self.axis, not self.inv).apply((g, inds.data))
        return gx, None


def permutate(x, indices, axis=0, inv=False):
    """Permutates a given variable along an axis.

    This function permutate ``x`` with given ``indices``.
    That means ``y[i] = x[indices[i]]`` for all ``i``.
    Note that this result is same as ``y = x.take(indices)``.
    ``indices`` must be a permutation of ``[0, 1, ..., len(x) - 1]``.

    When ``inv`` is ``True``, ``indices`` is treated as its inverse.
    That means ``y[indices[i]] = x[i]``.

    Args:
        x (:class:`~chainer.Variable` or :class:`numpy.ndarray` or \
        :class:`cupy.ndarray`):
            Variable to permutate.
            A :math:`(s_1, s_2, ..., s_N)` -shaped float array.
        indices (:class:`~chainer.Variable` or :class:`numpy.ndarray` or \
        :class:`cupy.ndarray`):
            Indices to extract from the variable. A one-dimensional int array.
        axis (int): Axis that the input array is permutate along.
        inv (bool): If ``True``, ``indices`` is treated as its inverse.

    Returns:
        ~chainer.Variable: Output variable.

    .. admonition:: Example

        >>> x = np.arange(6).reshape((3, 2)).astype(np.float32)
        >>> x
        array([[0., 1.],
               [2., 3.],
               [4., 5.]], dtype=float32)
        >>> indices = np.array([2, 0, 1], np.int32)
        >>> y = F.permutate(x, indices)
        >>> y.data
        array([[4., 5.],
               [0., 1.],
               [2., 3.]], dtype=float32)
        >>> y = F.permutate(x, indices, inv=True)
        >>> y.data
        array([[2., 3.],
               [4., 5.],
               [0., 1.]], dtype=float32)
        >>> indices = np.array([1, 0], np.int32)
        >>> y = F.permutate(x, indices, axis=1)
        >>> y.data
        array([[1., 0.],
               [3., 2.],
               [5., 4.]], dtype=float32)

    """
    y, = Permutate(axis, inv).apply((x, indices))
    return y

@@ -1036,36 +1036,68 @@
 __(self,
+ indices,
  axis
-=0
 , inv
-=False):
+):%0A        self.indices = indices
 %0A       
@@ -1228,18 +1228,8 @@
 'x',
- 'indices'
 ))%0A 
@@ -1242,25 +1242,16 @@
  x_type,
- ind_type
  = in_ty
@@ -1411,176 +1411,8 @@
 s)%0A%0A
-        type_check.expect(%0A            ind_type.dtype.kind == 'i',%0A            ind_type.ndim == 1,%0A            x_type.shape%5Bself.axis%5D == ind_type.shape%5B0%5D,%0A        )%0A%0A
     
@@ -1619,40 +1619,27 @@
     
-self.retain_inputs((1,))
+x, = inputs
 %0A       
  x, 
@@ -1638,19 +1638,16 @@
     
- x,
  inds = 
 inpu
@@ -1642,21 +1642,27 @@
  inds = 
-input
+self.indice
 s%0A%0A     
@@ -1832,29 +1832,25 @@
     
-inds = self.in
+g, = grad_out
 puts
-%5B1%5D
 %0A   
@@ -1854,32 +1854,34 @@
         
-g, = grad_output
+inds = self.indice
 s%0A      
@@ -1894,24 +1894,30 @@
 = Permutate(
+inds, 
 self.axis, n
@@ -1942,18 +1942,8 @@
 ((g,
- inds.data
 ))%0A 
@@ -1963,13 +1963,8 @@
  gx,
- None
 %0A%0A%0Ad
@@ -3785,16 +3785,236 @@
     %22%22%22%0A
+    if indices.dtype.kind != 'i' or indices.ndim != 1:%0A        raise ValueError(%0A            'indices should be a one-dimensional int array')%0A    if isinstance(indices, chainer.Variable):%0A        indices = indices.array%0A
     y, =
@@ -4024,16 +4024,25 @@
 rmutate(
+indices, 
 axis, in
@@ -4053,24 +4053,16 @@
 pply((x,
- indices
 ))%0A    r

Stop passing `indices` as an input variable