import chainer
from chainer import backend
import chainer.utils


class Send(chainer.Function):
    """Send elements to target process."""

    def __init__(self, comm, peer_rank, peer_tag):
        chainer.utils.experimental('chainermn.functions.Send')
        self.comm = comm
        self.peer_rank = peer_rank
        self.peer_tag = peer_tag

    @property
    def label(self):
        return "{} (peer_rank: {})".format(
            self.__class__.__name__,
            self.peer_rank)

    def forward(self, inputs):
        xp = backend.get_array_module(*inputs)

        # The last input is dummy variable, to retain gradient computation
        # of this function.
        xs = inputs[:-1]

        if len(xs) == 1:
            xs = xs[0]

        self.comm.send(xs, self.peer_rank, self.peer_tag)

        # Return an empty variable, which serves as "delegate_variable."
        return xp.array([], dtype=xp.float32),

    def backward(self, inputs, grad_outputs):
        xp = backend.get_array_module(*inputs)
        dummy_grad = xp.array([], dtype=xp.float32)
        grad = self.comm.recv(self.peer_rank, self.peer_tag)
        if isinstance(grad, tuple):
            return tuple([xp.array(gy) for gy in grad] + [dummy_grad])
        else:
            return xp.array(grad), dummy_grad


class Recv(chainer.Function):
    """Receive elements from target process."""

    def __init__(self, comm, peer_rank, peer_tag):
        chainer.utils.experimental('chainermn.functions.Recv')
        self.comm = comm
        self.peer_rank = peer_rank
        self.peer_tag = peer_tag

    def __call__(self, *inputs):
        xp = backend.get_array_module(*inputs)

        if inputs == ():
            # Expected to be invoked without any args in usual case.
            dummy_var = chainer.Variable(xp.array([], dtype=xp.float32))
            dummy_var.name = 'dummy_var'
            return super(Recv, self).__call__(dummy_var)

        else:
            # Used for retaining computational graph.
            return super(Recv, self).__call__(*inputs)

    @property
    def label(self):
        return "{} (peer_rank: {})".format(
            self.__class__.__name__,
            self.peer_rank)

    def forward(self, inputs):
        data = self.comm.recv(self.peer_rank, self.peer_tag)

        if not isinstance(data, tuple):
            data = tuple([data])

        return data

    def backward(self, inputs, grad_outputs):
        xp = backend.get_array_module(*inputs)
        self.comm.send(grad_outputs, self.peer_rank, self.peer_tag)

        # dummy_var is needed to maintain Chainer's constraint.
        if inputs == ():
            dummy_var = tuple([xp.array([], dtype=xp.float32)])
        else:
            dummy_var = tuple([xp.zeros(x.shape, dtype=xp.float32)
                               for x in inputs])

        return dummy_var


def send(x, communicator, rank, tag=0):
    """Send elements to target process.

    This function returns a dummy variable only holding the computational
    graph. If ``backward()`` is invoked by this dummy variable, it will
    try to receive gradients from the target process and send them back
    to the parent nodes.

    Args:
        x (Variable): Variable holding a matrix which you would like to send.
        communicator (chainer.communicators.CommunicatorBase):
            ChainerMN communicator.
        rank (int): Target process specifier.
        tag (int): Optional message ID (MPI feature).

    Returns:
        ~chainer.Variable:
            A dummy variable with no actual data, only holding the
            computational graph. Please refer
            ``chainermn.functions.pseudo_connect`` for detail.

    """
    chainer.utils.experimental('chainermn.functions.send')

    if rank == communicator.rank:
        raise ValueError(
            'rank must be different from communicator rank, '
            'otherwise deadlock occurs')

    xp = backend.get_array_module(*x)

    # Dummy variable to retain gradient computation of send,
    # otherwise the corresponding recv will cause deadlock in backward
    # in the case where all inputs for this function does not require_grad.
    dummy_var = chainer.Variable(xp.array([], dtype=xp.float32))

    if isinstance(x, list) or isinstance(x, tuple):
        inputs = x + type(x)([dummy_var])
        delegate_variable = Send(
            communicator, peer_rank=rank, peer_tag=tag)(*inputs)
    else:
        delegate_variable = Send(
            communicator, peer_rank=rank, peer_tag=tag)(x, dummy_var)

    delegate_variable.name = 'delegate_variable'
    return delegate_variable


def recv(communicator, rank, delegate_variable=None, tag=0, force_tuple=False):
    """Receive elements from target process.

    This function returns data received from target process. If ``backward()``
    is invoked, it will try to send gradients to the target process.
    The received array will be on the current CUDA device if the corresponding
    ``send()`` is invoked with arrays on GPU.
    Please be aware that the current CUDA device is intended one.
    (``https://docs-cupy.chainer.org/en/stable/tutorial/basic.html#current-device``)

    .. note::
        If you define non-connected computational graph on one process,
        you have to use ``delegate_variable`` to specify the output of
        previous computational graph component.
        Otherwise ``backward()`` does not work well.
        Please refer ``chainermn.functions.pseudo_connect`` for detail.

    Args:
        communicator (chainer.communicators.CommunicatorBase):
            ChainerMN communicator.
        rank (int): Target process specifier.
        delegate_variable (chainer.Variable):
            Pointer to the other non-connected component.
        tag (int): Optional message ID (MPI feature).
        force_tuple (bool): If ``False`` (the default) a Variable will be
            returned when the number of outputs is one. Otherwise, this
            method returns a tuple even when the number of outputs is one.

    Returns:
        ~chainer.Variable:
            Data received from target process. If ``backward()`` is invoked
            by this variable, it will send gradients to the target process.

    """
    chainer.utils.experimental('chainermn.functions.recv')

    if rank == communicator.rank:
        raise ValueError(
            'rank must be different from communicator rank, '
            'otherwise deadlock occurs')

    if delegate_variable is None:
        res = Recv(
            communicator,
            peer_rank=rank,
            peer_tag=tag)()
    else:
        delegate_variable.name = 'delegate_variable'
        res = Recv(
            communicator,
            peer_rank=rank,
            peer_tag=tag)(delegate_variable)

    if force_tuple and not isinstance(res, tuple):
        return tuple([res])
    else:
        return res
