def _MPIVisualizer_visualize(d, U, **kwargs):
    d = mpi.get_object(d)
    if isinstance(U, tuple):
        U = tuple(mpi.get_object(u) for u in U)
    else:
        U = mpi.get_object(U)
    d.visualize(U, **kwargs)

def _MPIVectorAutoComm_dot(self, other):
    self = mpi.get_object(self)
    other = mpi.get_object(other)
    local_result = self.dot(other)
    assert local_result.dtype == np.float64
    results = np.empty((mpi.size,), dtype=np.float64) if mpi.rank0 else None
    mpi.comm.Gather(local_result, results, root=0)
    if mpi.rank0:
        return np.sum(results)

def _MPIVectorArrayAutoComm_pairwise_dot(self, other, ind=None, o_ind=None):
    self = mpi.get_object(self)
    other = mpi.get_object(other)
    local_results = self.pairwise_dot(other, ind=ind, o_ind=o_ind)
    assert local_results.dtype == np.float64
    results = np.empty((mpi.size,) + local_results.shape, dtype=np.float64) if mpi.rank0 else None
    mpi.comm.Gather(local_results, results, root=0)
    if mpi.rank0:
        return np.sum(results, axis=0)

 def __init__(self, obj_id, functional=False, vector=False, with_apply2=False,
              pickle_subtypes=True, array_type=MPIVectorArray):
     assert not (functional and vector)
     self.obj_id = obj_id
     self.op = op = mpi.get_object(obj_id)
     self.functional = functional
     self.vector = vector
     self.with_apply2 = with_apply2
     self.pickle_subtypes = pickle_subtypes
     self.array_type = array_type
     self.linear = op.linear
     self.name = op.name
     self.build_parameter_type(inherits=(op,))
     if vector:
         self.source = NumpyVectorSpace(1)
         assert self.source == op.source
     else:
         subtypes = mpi.call(_MPIOperator_get_source_subtypes, obj_id, pickle_subtypes)
         if all(subtype == subtypes[0] for subtype in subtypes):
             subtypes = (subtypes[0],)
         self.source = VectorSpace(array_type, (op.source.type, subtypes))
     if functional:
         self.range = NumpyVectorSpace(1)
         assert self.range == op.range
     else:
         subtypes = mpi.call(_MPIOperator_get_range_subtypes, obj_id, pickle_subtypes)
         if all(subtype == subtypes[0] for subtype in subtypes):
             subtypes = (subtypes[0],)
         self.range = VectorSpace(array_type, (op.range.type, subtypes))

 def __init__(self, obj_id, mpi_range, mpi_source, with_apply2=False, pickle_local_spaces=True,
              space_type=MPIVectorSpace):
     assert mpi_source or mpi_range
     self.obj_id = obj_id
     self.mpi_source = mpi_source
     self.mpi_range = mpi_range
     self.op = op = mpi.get_object(obj_id)
     self.with_apply2 = with_apply2
     self.pickle_local_spaces = pickle_local_spaces
     self.space_type = space_type
     self.linear = op.linear
     self.name = op.name
     self.build_parameter_type(op)
     if mpi_source:
         local_spaces = mpi.call(_MPIOperator_get_local_spaces, obj_id, True, pickle_local_spaces)
         if all(ls == local_spaces[0] for ls in local_spaces):
             local_spaces = (local_spaces[0],)
         self.source = space_type(local_spaces)
     else:
         self.source = op.source
     if mpi_range:
         local_spaces = mpi.call(_MPIOperator_get_local_spaces, obj_id, False, pickle_local_spaces)
         if all(ls == local_spaces[0] for ls in local_spaces):
             local_spaces = (local_spaces[0],)
         self.range = space_type(local_spaces)
     else:
         self.range = op.range

def mpi_wrap_operator(obj_id, mpi_range, mpi_source, with_apply2=False, pickle_local_spaces=True,
                      space_type=MPIVectorSpace):
    """Wrap MPI distributed local |Operators| to a global |Operator| on rank 0.

    Given MPI distributed local |Operators| referred to by the
    :class:`~pymor.tools.mpi.ObjectId` `obj_id`, return a new |Operator|
    which manages these distributed operators from rank 0. This
    is done by instantiating :class:`MPIOperator`. Additionally, the
    structure of the wrapped operators is preserved. E.g. |LincombOperators|
    will be wrapped as a |LincombOperator| of :class:`MPIOperators <MPIOperator>`.

    Parameters
    ----------
    See :class:`MPIOperator`.

    Returns
    -------
    The wrapped |Operator|.
    """
    op = mpi.get_object(obj_id)
    if isinstance(op, LincombOperator):
        obj_ids = mpi.call(_mpi_wrap_operator_LincombOperator_manage_operators, obj_id)
        return LincombOperator([mpi_wrap_operator(o, mpi_range, mpi_source, with_apply2, pickle_local_spaces,
                                                  space_type)
                                for o in obj_ids], op.coefficients, name=op.name)
    elif isinstance(op, VectorArrayOperator):
        array_obj_id, local_spaces = mpi.call(_mpi_wrap_operator_VectorArrayOperator_manage_array,
                                              obj_id, pickle_local_spaces)
        if all(ls == local_spaces[0] for ls in local_spaces):
            local_spaces = (local_spaces[0],)
        return VectorArrayOperator(space_type(local_spaces).make_array(array_obj_id),
                                   adjoint=op.adjoint, name=op.name)
    else:
        return MPIOperator(obj_id, mpi_range, mpi_source, with_apply2, pickle_local_spaces, space_type)

def _mpi_wrap_operator_VectorArrayOperator_manage_array(obj_id):
    op = mpi.get_object(obj_id)
    array_obj_id = mpi.manage_object(op._array)
    subtypes = mpi.comm.gather(op._array.subtype, root=0)
    mpi.remove_object(obj_id)
    if mpi.rank0:
        return array_obj_id, tuple(subtypes)

def mpi_wrap_operator(obj_id, functional=False, vector=False, with_apply2=False,
                      pickle_subtypes=True, array_type=MPIVectorArray):
    """Wrap MPI distributed local |Operators| to a global |Operator| on rank 0.

    Given MPI distributed local |Operators| referred to by the
    `~pymor.tools.mpi.ObjectId` `obj_id`, return a new |Operator|
    which manages these distributed operators from rank 0. This
    is done by instantiating :class:`MPIOperator`. Additionally, the
    structure of the wrapped operators is preserved. E.g. |LincombOperators|
    will be wrapped as a |LincombOperator| of :class:`MPIOperators`.

    Parameters
    ----------
    See :class:`MPIOperator`.

    Returns
    -------
    The wrapped |Operator|.
    """
    op = mpi.get_object(obj_id)
    if isinstance(op, LincombOperator):
        obj_ids = mpi.call(_mpi_wrap_operator_LincombOperator_manage_operators, obj_id)
        return LincombOperator([mpi_wrap_operator(o, functional, vector, with_apply2, pickle_subtypes, array_type)
                                for o in obj_ids], op.coefficients, name=op.name)
    elif isinstance(op, VectorArrayOperator):
        array_obj_id, subtypes = mpi.call(_mpi_wrap_operator_VectorArrayOperator_manage_array, obj_id, pickle_subtypes)
        if all(subtype == subtypes[0] for subtype in subtypes):
            subtypes = (subtypes[0],)
        return VectorArrayOperator(array_type(type(op._array), subtypes, array_obj_id),
                                   transposed=op.transposed, name=op.name)
    else:
        return MPIOperator(obj_id, functional, vector, with_apply2, pickle_subtypes, array_type)

def mpi_wrap_discretization(obj_id, use_with=False, with_apply2=False, array_type=MPIVectorArray):
    """Wrap MPI distributed local |Discretizations| to a global |Discretization| on rank 0.

    Given MPI distributed local |Discretizations| referred to by the
    `~pymor.tools.mpi.ObjectId` `obj_id`, return a new |Discretization|
    which manages these distributed discretizations from rank 0. This
    is done by first wrapping all |Operators| of the |Discretization| using
    :func:`~pymor.operators.mpi.mpi_wrap_operator`.

    When `use_with` is `False`, an :class:`MPIDiscretization` is instatiated
    with the wrapped operators. A call to
    :meth:`~pymor.discretizations.interfaces.DiscretizationInterface.solve`
    will then use an MPI parallel call to the
    :meth:`~pymor.discretizations.interfaces.DiscretizationInterface.solve`
    methods of the wrapped local |Discretizations| to obtain the solution.
    This is usually what you want when the actual solve is performed by
    an implementation in the external solver.

    When `use_with` is `True`, :meth:`~pymor.core.interfaces.ImmutableInterface.with_`
    is called on the local |Discretization| on rank 0, to obtain a new
    |Discretization| with the wrapped MPI |Operators|. This is mainly useful
    when the local discretizations are generic |Discretizations| as in
    :mod:`pymor.discretizations.basic` and
    :meth:`~pymor.discretizations.interfaces.DiscretizationInterface.solve`
    is implemented directly in pyMOR via operations on the contained
    |Operators|.

    Parameters
    ----------
    obj_id
        :class:`~pymor.tools.mpi.ObjectId` of the local |Discretization|
        on each rank.
    use_with
        See above.
    with_apply2
        See :class:`~pymor.operators.mpi.MPIOperator`.
    array_type
        See :class:`~pymor.operators.mpi.MPIOperator`.
    """

    operators, functionals, vectors, products = \
        mpi.call(_mpi_wrap_discretization_manage_operators, obj_id)

    operators = {k: mpi_wrap_operator(v, with_apply2=with_apply2, array_type=array_type) if v else None
                 for k, v in operators.iteritems()}
    functionals = {k: mpi_wrap_operator(v, functional=True, with_apply2=with_apply2, array_type=array_type) if v else None
                   for k, v in functionals.iteritems()}
    vectors = {k: mpi_wrap_operator(v, vector=True, with_apply2=with_apply2, array_type=array_type) if v else None
               for k, v in vectors.iteritems()}
    products = {k: mpi_wrap_operator(v, with_apply2=with_apply2, array_type=array_type) if v else None
                for k, v in products.iteritems()} if products else None

    if use_with:
        d = mpi.get_object(obj_id)
        visualizer = MPIVisualizer(obj_id)
        return d.with_(operators=operators, functionals=functionals, vector_operators=vectors, products=products,
                       visualizer=visualizer, cache_region=None)
    else:
        return MPIDiscretization(obj_id, operators, functionals, vectors, products, array_type=array_type)

def _MPIDiscretization_get_subtypes(self, pickle_subtypes):
    self = mpi.get_object(self)
    subtype = self.solution_space.subtype
    if not pickle_subtypes:
        subtype = _register_subtype(subtype)
    subtypes = mpi.comm.gather(subtype, root=0)
    if mpi.rank0:
        return tuple(subtypes)

def _mpi_wrap_discretization_manage_operators(obj_id):
    d = mpi.get_object(obj_id)
    operators = {k: mpi.manage_object(v) if v else None for k, v in sorted(d.operators.iteritems())}
    functionals = {k: mpi.manage_object(v) if v else None for k, v in sorted(d.functionals.iteritems())}
    vectors = {k: mpi.manage_object(v) if v else None for k, v in sorted(d.vector_operators.iteritems())}
    products = {k: mpi.manage_object(v) if v else None for k, v in sorted(d.products.iteritems())} if d.products else None
    if mpi.rank0:
        return operators, functionals, vectors, products

def _MPIOperator_get_local_spaces(self, source, pickle_local_spaces):
    self = mpi.get_object(self)
    local_space = self.source if source else self.range
    if not pickle_local_spaces:
        local_space = _register_local_space(local_space)
    local_spaces = mpi.comm.gather(local_space, root=0)
    if mpi.rank0:
        return tuple(local_spaces)

def _MPIDiscretization_get_local_spaces(self, pickle_local_spaces):
    self = mpi.get_object(self)
    local_space = self.solution_space
    if not pickle_local_spaces:
        local_space = _register_local_space(local_space)
    local_spaces = mpi.comm.gather(local_space, root=0)
    if mpi.rank0:
        return tuple(local_spaces)

def _MPIOperator_get_range_subtypes(self, pickle_subtypes):
    self = mpi.get_object(self)
    subtype = self.range.subtype
    if not pickle_subtypes:
        subtype = _register_subtype(subtype)
    subtypes = mpi.comm.gather(subtype, root=0)
    if mpi.rank0:
        return tuple(subtypes)

 def _map(self, function, chunks, **kwargs):
     payload = mpi.get_object(self._payload)
     payload[0] = chunks
     try:
         result = mpi.call(mpi.function_call, _worker_map_function, self._payload, function, **kwargs)
     finally:
         payload[0] = None
     return result

 def _apply_only(self, function, worker, *args, **kwargs):
     payload = mpi.get_object(self._payload)
     payload[0] = (function, args, kwargs)
     try:
         result = mpi.call(mpi.function_call, _single_worker_call_function, self._payload, worker)
     finally:
         payload[0] = None
     return result

def _MPIVectorArrayAutoComm_l2_norm2(self, ind=None):
    self = mpi.get_object(self)
    local_results = self.l2_norm2(ind=ind)
    assert local_results.dtype == np.float64
    results = np.empty((mpi.size,) + local_results.shape, dtype=np.float64) if mpi.rank0 else None
    mpi.comm.Gather(local_results, results, root=0)
    if mpi.rank0:
        return np.sum(results, axis=0)

def _MPIVectorAutoComm_l2_norm(self):
    self = mpi.get_object(self)
    local_result = self.l2_norm2()
    assert local_result.dtype == np.float64
    results = np.empty((mpi.size,), dtype=np.float64) if mpi.rank0 else None
    mpi.comm.Gather(local_result, results, root=0)
    if mpi.rank0:
        return np.sqrt(np.sum(results))

def _mpi_wrap_operator_VectorArrayOperator_manage_array(obj_id, pickle_local_spaces):
    op = mpi.get_object(obj_id)
    array_obj_id = mpi.manage_object(op._array)
    local_space = op._array.space
    if not pickle_local_spaces:
        local_space = _register_local_space(local_space)
    local_spaces = mpi.comm.gather(local_space, root=0)
    mpi.remove_object(obj_id)
    if mpi.rank0:
        return array_obj_id, tuple(local_spaces)

def _mpi_wrap_operator_VectorArrayOperator_manage_array(obj_id, pickle_subtypes):
    op = mpi.get_object(obj_id)
    array_obj_id = mpi.manage_object(op._array)
    subtype = op._array.subtype
    if not pickle_subtypes:
        subtype = _register_subtype(subtype)
    subtypes = mpi.comm.gather(subtype, root=0)
    mpi.remove_object(obj_id)
    if mpi.rank0:
        return array_obj_id, tuple(subtypes)