from six.moves import queue

import chainer
import chainermn
import chainermn.communicators
import chainermn.functions


class MultiNodeChainList(chainer.ChainList):
    """Combining multiple non-connected components of computational graph.

    This class combines each ``chainer.Chain``, which represents one of the
    non-connected component in compuational graph. In ``__call__()``,
    the returned object of ``chainer.Chain`` (which represents pointer)
    are passed to the next ``chainer.Chain``, in order to retain the
    computational graph connected and make backprop work properly.

    Users add each ``chainer.Chain`` by ``add_link()`` method. Each chain
    is invoked in forward computation according to the order they are added,
    and in backward computation according to the reversed order.

    .. admonition:: Example (basic usage)

        This is a simple example of the model which sends its outputs to
        rank=1 machine::

            import chainer
            import chainer.functions as F
            import chainermn


            class SimpleModelSub(chainer.Chain):

                def __init__(self, n_in, n_hidden, n_out):
                    super(SimpleModelSub, self).__init__(
                        l1=L.Linear(n_in, n_hidden),
                        l2=L.Linear(n_hidden, n_out))

                def __call__(self, x):
                    h1 = F.relu(self.l1(x))
                    return self.l2(h1)


            class SimpleModel(chainermn.MultiNodeChainList):

                def __init__(self, comm, n_in, n_hidden, n_out):
                    super(SimpleModel, self).__init__(comm)
                    self.add_link(
                        SimpleModelSub(n_in, n_hidden, n_out),
                        rank_in=None,
                        rank_out=1)

    .. admonition:: Example (split MLP on 2 processes)

        This is the other example of two models interacting each other::

            import chainer
            import chainer.functions as F
            import chainermn


            class MLP(chainer.Chain):

                def __init__(self, n_in, n_hidden, n_out):
                    super(MLP, self).__init__(
                        l1=L.Linear(n_in, n_hidden),
                        l2=L.Linear(n_hidden, n_hidden),
                        l3=L.Linear(n_hidden, n_out))

                def __call__(self, x):
                    h1 = F.relu(self.l1(x))
                    h2 = F.relu(self.l2(h1))
                    return self.l3(h2)


            class Model0(chainermn.MultiNodeChainList):

                def __init__(self, comm):
                    super(Model0, self).__init__(comm)
                    self.add_link(
                        MLP(10000, 5000, 2000),
                        rank_in=None,
                        rank_out=1)
                    self.add_link(
                        MLP(100, 50, 10),
                        rank_in=1,
                        rank_out=None)


            class Model1(chainermn.MultiNodeChainList):

                def __init__(self, comm):
                    super(Model1, self).__init__(comm)
                    self.add_link(MLP(2000, 500, 100), rank_in=0, rank_out=0)


        ``Model0`` is expected to be on rank=0, and ``Model1`` is expected to
        be on rank=1. The first ``MLP`` in ``Model0`` will send its outputs
        to ``Model1``, then ``MLP`` in ``Model1`` will receive it and send
        its outputs to the second ``MLP`` in ``Model0``.

    .. admonition:: Example (sending tuples)

        This is the example for sending a tuple::

            import chainer
            import chainer.functions as F
            import chainermn

            class NN0(chainer.Chain):
                def __call__(self, x):
                    y0 = some_calculation_nn0_0(x)
                    y1 = some_calculation_nn1_1(x)
                    return y0, y1

            class NN1(chainer.Chain):
                def __call__(self, y):
                    y0, y1 = y  # unpack tuple from NN0
                    return some_calculation_nn1(y0, y1)

            class Model_on_Process_0(chainermn.MultiNodeChainList):
                def __init__(self, comm):
                    super(Model_on_Process_0, self).__init__(comm=comm)
                    self.add_link(NN0(), rank_in=None, rank_out=1)

            class Model_on_Process_1(chainermn.MultiNodeChainList):
                def __init__(self, comm):
                    super(Model_on_Process_1, self).__init__(comm=comm)
                    self.add_link(NN1(), rank_in=0, rank_out=None)

        In this example, ``Model_on_Process_0`` sends two elemental tuple
        ``(y0, y1)`` (returned by ``NN0.__call__``) to ``Model_on_Process_1``,
        which can be unpacked as shown in ``NN1.__call__``.

    Args:
        comm (chainermn.communicators._base.CommunicatorBase):
            ChainerMN communicator.
    """

    def __init__(self, comm):
        chainer.utils.experimental('chainermn.MultiNodeChainList')
        super(MultiNodeChainList, self).__init__()
        self._comm = comm
        self._rank_inouts = []

    def add_link(self, link, rank_in=None, rank_out=None):
        """Register one connected link with its inout rank.

        Args:
            link (chainer.Link): The link object to be registered.
            rank_in (int, list, or None):
                Ranks from which it receives data. If None is specified,
                the model does not receive from any machines.
            rank_out (int, list, or None):
                Ranks to which it sends data. If None is specified,
                the model will not send to any machine.
        """
        super(MultiNodeChainList, self).add_link(link)
        if isinstance(rank_in, int):
            rank_in = [rank_in]
        if isinstance(rank_out, int):
            rank_out = [rank_out]

        if rank_out is None:
            for _, _rank_out in self._rank_inouts:
                if _rank_out is None:
                    raise ValueError(
                        'MultiNodeChainList cannot have more than two '
                        'computational graph component whose rank_out is None')

        self._rank_inouts.append((rank_in, rank_out))

    def __call__(self, *inputs):
        comm_queue = queue.Queue()
        y = None
        delegate_variable = None

        for i_comp, (f, (rank_in, rank_out)) in \
                enumerate(zip(self._children, self._rank_inouts)):
            x = None

            if rank_in is None:  # Use inputs.
                if i_comp == 0:
                    x = f(*inputs)
                else:
                    # If the graph component is not the first one,
                    # backprop to the previous graph component must be
                    # guaranteed.
                    x = chainermn.functions.pseudo_connect(
                        delegate_variable,
                        *inputs)
                    x = f(x)

            else:  # Receive inputs from the other machines.
                # Preprocess: receiving inputs from the other machines.
                xs = []
                for _rank_in in rank_in:
                    if _rank_in == self._comm.rank:
                        # Receive inputs from itself.
                        if delegate_variable is None:
                            _x = comm_queue.get()
                        else:
                            _x = chainermn.functions.pseudo_connect(
                                delegate_variable,
                                comm_queue.get())
                    else:
                        _x = chainermn.functions.recv(
                            self._comm,
                            rank=_rank_in,
                            delegate_variable=delegate_variable)

                    xs.append(_x)

                    # Guarantee the backward path to the previous graph
                    # component to be executed in the last to avoid dead-lock.
                    delegate_variable = _x

                # Guarantee backprop on the same edge exactly once.
                delegate_variable = None

                # Actual forward.
                x = f(*tuple(xs))

            if rank_out is None:  # Return outputs.
                assert y is None, "MultiNodeChainList cannot have more than "\
                    "two computational graph component whose rank_out is None"
                y = x  # model output
                delegate_variable = y

            else:  # Send outputs to the other machines.
                for i_comp, _rank_out in enumerate(rank_out):
                    if _rank_out == self._comm.rank:
                        # Send outputs to itself.
                        if delegate_variable is not None:
                            x = chainermn.functions.pseudo_connect(
                                delegate_variable,
                                x)
                        comm_queue.put(x)
                        delegate_variable = x
                    elif i_comp == 0:
                        delegate_variable = chainermn.functions.send(
                            x, self._comm,
                            rank=_rank_out)
                    else:
                        # If the model has multiple targets for send,
                        # we must guarantee backwards of each send to be
                        # called in the reversed order.
                        if delegate_variable is not None:
                            x = chainermn.functions.pseudo_connect(
                                delegate_variable,
                                x)
                        delegate_variable = chainermn.functions.send(
                            x, self._comm,
                            rank=_rank_out)

        if not comm_queue.empty():
            raise ValueError(
                'Communication queue is not empty at the end of forward.'
                'Make sure if all rank_in and rank_out correspond each other.')

        # Return.
        if y is delegate_variable:
            # The last computational graph component returns model output.
            return y
        elif y is not None:
            # The intermediate graph component returns model output.
            return chainermn.functions.pseudo_connect(delegate_variable, y)
        else:
            # Do not have any model output.
            return delegate_variable
