def wrap(ag, compound='atoms'):
    """
    Shift the contents of a given AtomGroup back into the unit cell. ::

       +-----------+          +-----------+
       |           |          |           |
       |         3 | 6        | 6       3 |
       |         ! | !        | !       ! |
       |       1-2-|-5-8  ->  |-5-8   1-2-|
       |         ! | !        | !       ! |
       |         4 | 7        | 7       4 |
       |           |          |           |
       +-----------+          +-----------+

    Example
    -------

    .. code-block:: python

        ag = u.atoms
        transform = mda.transformations.wrap(ag)
        u.trajectory.add_transformations(transform)

    Parameters
    ----------

    ag: Atomgroup
        Atomgroup to be wrapped in the unit cell
    compound : {'atoms', 'group', 'residues', 'segments', 'fragments'}, optional
        The group which will be kept together through the shifting process.

    Notes
    -----
    When specifying a `compound`, the translation is calculated based on
    each compound. The same translation is applied to all atoms
    within this compound, meaning it will not be broken by the shift.
    This might however mean that not all atoms from the compound are
    inside the unit cell, but rather the center of the compound is.

    Returns
    -------
    MDAnalysis.coordinates.base.Timestep

    """

    def wrapped(ts):
        ag.wrap(compound=compound)

        return ts

    return wrapped