<commit_before># Copyright 2014, <NAME> <<EMAIL>>
#
# This file is part of Azrael (https://github.com/olitheolix/azrael)
#
# Azrael is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# Azrael is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with Azrael. If not, see <http://www.gnu.org/licenses/>.

"""
Physics manager.
"""
import os
import sys
import zmq
import time
import signal
import pickle
import itertools
import numpy as np

import azrael.igor
import azrael.database
import azrael.vectorgrid
import azrael.bullet_api
import azrael.util as util
import azrael.config as config
import azrael.rb_state as rb_state
import azrael.leo_api as leoAPI

from IPython import embed as ipshell
from azrael.types import _RigidBodyState
from azrael.types import typecheck, RetVal, WPData, WPMeta, Forces

# Convenience.
RigidBodyState = rb_state.RigidBodyState
RigidBodyStateOverride = rb_state.RigidBodyStateOverride


@typecheck
def sweeping(data: list, labels: np.ndarray, dim: str):
    """
    Return sets of overlapping AABBs in the dimension ``dim``.

    This function implements the 'Sweeping' algorithm to determine which sets
    of AABBs overlap.

    Sweeping is straightforward: sort all start/stop positions and determine
    the overlapping sets.

    The returned sets do not contain the ``data`` elements but their
    corresponding ``labels`` to be more memory efficient.

    The `data` is a list of dictionaries that denotes the half widths of the
    AABBs in the respective direction:
      data = [{'x': [x0, x1], 'y': [y0, y1], 'z': [z0, z1]},
              {'x': [x0, x1], 'y': [y0, y1], 'z': [z0, z1]},
              ...]
    The elements denote the AABB of the corresponding `label` at the same
    index. The dictionary keys 'x', 'y', and 'z' are (more or less) hard coded,
    and `dim` must refer to one of them (eg dim = 'y' to find the sets that
    overlap in the 'y' dimension).

    :param list[dict] data: list of dictionaries containing the AABB boundaries
       for 'x', 'y', and 'z'.
    :param np.int64 labels: integer array to label the elements in data.
    :param str dim: the axis to check (must be one of ['x', 'y', 'z'])
    :return: list of sets. Each set contains elements from ``labels``.
    :rtype: list of sets
    """
    assert len(labels) == len(data)

    # Convenience.
    N = 2 * len(data)

    # Pre-allocate arrays for start/stop position, objID, and an
    # increment/decrement array used for convenient processing afterwards.
    arr_pos = np.zeros(N, np.float64)
    arr_lab = np.zeros(N, np.int64)
    arr_inc = np.zeros(N, np.int8)

    # Fill the arrays.
    for ii in range(len(data)):
        arr_pos[2 * ii: 2 * ii + 2] = np.array(data[ii][dim])
        arr_lab[2 * ii: 2 * ii + 2] = labels[ii]
        arr_inc[2 * ii: 2 * ii + 2] = [+1, -1]

    # Sort all three arrays according to the start/stop positions.
    idx = np.argsort(arr_pos)
    arr_lab = arr_lab[idx]
    arr_inc = arr_inc[idx]

    # Output array.
    out = []

    # Sweep over the sorted data and compile the list of object sets.
    sumVal = 0
    setObjs = set()
    for (inc, objID) in zip(arr_inc, arr_lab):
        # Update the index variable and add the current object to the set.
        sumVal += inc
        setObjs.add(objID)

        # A new set of overlapping AABBs is complete whenever `sumVal`
        # reaches zero.
        if sumVal == 0:
            out.append(setObjs)
            setObjs = set()

        # Safety check: sumVal can never be negative.
        assert sumVal >= 0
    return RetVal(True, None, out)


@typecheck
def computeCollisionSetsAABB(SVs: dict, AABBs: dict):
    """
    Return broadphase collision sets for the objects in ``SVs``.

    :param dict SVs: Dictionary of Body States.
    :param dict AABBs: Dictionary of AABBs.
    :return: each list contains a unique set of overlapping objects.
    :rtype: list of lists
    """
    # Sanity check.
    if set(SVs.keys()) != set(AABBs.keys()):
        return RetVal(False, 'SVs and AABBs are inconsisten', None)

    # The 'sweeping' function requires a list of dictionaries. Each dictionary
    # must contain the min/max spatial extent in x/y/z direction.
    data = []
    IDs = list(SVs.keys())
    for objID in IDs:
        if (SVs[objID] is None) or (AABBs[objID] is None):
            continue
        sv, aabb = SVs[objID], AABBs[objID]
        pos = sv.position
        x0, x1 = pos[0] - aabb, pos[0] + aabb
        y0, y1 = pos[1] - aabb, pos[1] + aabb
        z0, z1 = pos[2] - aabb, pos[2] + aabb

        data.append({'x': [x0, x1], 'y': [y0, y1], 'z': [z0, z1]})
    del SVs, AABBs

    # Enumerate the objects.
    labels = np.arange(len(IDs))

    # Determine the sets of objects that overlap 'x' direction.
    stage_0 = sweeping(data, labels, 'x').data

    # Iterate over all the just found sets. For each, determine the sets that
    # overlap in the 'y' dimension.
    stage_1 = []
    for subset in stage_0:
        tmpData = [data[_] for _ in subset]
        tmpLabels = np.array(tuple(subset), np.int64)
        stage_1.extend(sweeping(tmpData, tmpLabels, 'y').data)

    # Iterate over all the sets that overlap in 'x' and 'y' dimension. For
    # each, determine which also overalp in the 'z' dimension.
    stage_2 = []
    for subset in stage_1:
        tmpData = [data[_] for _ in subset]
        tmpLabels = np.array(tuple(subset), np.int64)
        stage_2.extend(sweeping(tmpData, tmpLabels, 'z').data)

    # Convert the labels back to object IDs.
    out = [[IDs[objID] for objID in objIDs] for objIDs in stage_2]

    # Return the sets.
    return RetVal(True, None, out)


def mergeConstraintSets(constraintPairs: tuple, collSets: (tuple, list)):
    """
    Merge all the sets in ``collSets`` that contain any of the
    ``constraintPairs``.

    The purpose of this function is to merge those collision sets that are
    connected via a constraint. Typically, this function takes the output of
    ``computeCollisionSets`` as the ``collSets`` argument.

    :param list constraintPairs: list of 2-tuples eg [(1, 2), (1, 5), ...].
    :param list collSets: a set/list of sets/lists.
    :return: list of collision sets.
    """
    # Merge the collision sets that are linked via constraints.
    for (a, b) in constraintPairs:
        # Find (and remove) the set(s) that contain object_a or object_b (or
        # both).
        s_a = [collSets.pop(ii) for ii, v in enumerate(collSets) if a in v]
        s_b = [collSets.pop(ii) for ii, v in enumerate(collSets) if b in v]

        # Sanity check: each element can be in at most one set.
        assert 0 <= len(s_a) < 2
        assert 0 <= len(s_b) < 2

        # Remove the list from the list comprehension.
        s_a = s_a[0] if len(s_a) == 1 else []
        s_b = s_b[0] if len(s_b) == 1 else []

        # Merge the collision sets.
        collSets.append(s_a + s_b)
    return RetVal(True, None, collSets)


def getFinalCollisionSets(
        constraintPairs: list, allObjects: list, allAABBs: list):
    """
    Return the collision sets.

    This function calls the broadphase solver and then merges those collision
    sets that are connected via a constraint.

    .. note:: this function is tightly coupled to the inner working of Leonard.
              Possibly it should not be a standalone function.

    :param list constraintPairs: list of 2-tuples eg [(1, 2), (1, 5), ...].
    :param list allObjects: Leonard's object cache.
    :param list allAABBs: Leonard's AABB cache.
    :return: list of non-overlapping collision sets.
    """
    # Broadphase based on AABB only.
    ret = computeCollisionSetsAABB(allObjects, allAABBs)
    if not ret.ok:
        msg = 'ComputeCollisionSetsAABB returned an error'
        self.logit.error(msg)
        return RetVal(False, msg, None)

    # Merge all collision sets that have objects which are connected by a
    # constraint.
    collSets = ret.data
    ret = mergeConstraintSets(constraintPairs, collSets)
    if not ret.ok:
        msg = 'mergeConstraintSets returned an error'
        self.logit.error(msg)
        return RetVal(False, msg, None)
    return ret


class LeonardBase(config.AzraelProcess):
    """
    Base class for Physics manager.

    No physics is actually computed here. The class serves mostly as an
    interface for the actual Leonard implementations, as well as a test
    framework.
    """
    def __init__(self):
        super().__init__()

        # Create the DB handles.
        self._DB_SV = azrael.database.dbHandles['SV']

        # Create an Igor instance.
        self.igor = azrael.igor.Igor()

        self.allObjects = {}
        self.allAABBs = {}
        self.allForces = {}

    def setup(self):
        """
        Stub for initialisation code that cannot go into the constructor.

        This is typically necessary for code that must not execute until
        *after* the class forked into a new process. ZeroMQ contexts are a good
        example for this - *never* fork processes *after* creating the context.
        """
        pass

    def getGridForces(self, idPos: dict):
        """
        Return dictionary of force values for every object in ``idPos``.

        The ``idPos`` argument is a {objID_1: sv_1, objID_2, sv_2, ...}
        dictionary.

        The returned dictionary has the same keys as ``idPos``.

        :param dict idPos: dictionary with objIDs and corresponding SVs.
        :return dict: {objID_k: force_k}
        """
        # Convenience.
        vg = azrael.vectorgrid

        # Extract the keys and values in the same order.
        objIDs = list(idPos.keys())
        positions = [idPos[_] for _ in objIDs]

        # Query the grid values at all positions.
        ret = vg.getValues('force', positions)
        if not ret.ok:
            return RetVal(False, ret.msg, None)

        # Overwrite the default values.
        gridForces = {objID: val for objID, val in zip(objIDs, ret.data)}
        return RetVal(True, None, gridForces)

    @typecheck
    def totalForceAndTorque(self, objID: int):
        """
        Return the total force- and torque on the object.

        The returned values are the sum of all booster forces (correctly
        oriented relative to the object) and the force a user may have specifed
        directly.

        Note that this function does not account for the forces from the 'force
        grid'.

        :param int objID: return the force and torque for this object.
        :return: the force and torque as two Python lists (not NumPy arrays)
        :rtype: (list, list)
        """
        # Convenience.
        sv = self.allObjects[objID]
        f = self.allForces[objID]

        # Construct the Quaternion of the object based on its orientation.
        quat = util.Quaternion(sv.orientation[3], sv.orientation[:3])

        # Fetch the force vector for the current object from the DB.
        force = np.array(f.forceDirect, np.float64)
        torque = np.array(f.torqueDirect, np.float64)

        # Add the booster's contribution to force and torque.
        # Note: We cannot do this directly since the booster force/torque were
        # specified in object coordinates. We thus rotate them to world
        # coordinates before adding them to the total force.
        force += quat * f.forceBoost
        torque += quat * f.torqueBoost

        # Convert to Python lists.
        return force.tolist(), torque.tolist()

    @typecheck
    def step(self, dt: (int, float), maxsteps: int):
        """
        Advance the simulation by ``dt`` using at most ``maxsteps``.

        This method will use a primitive Euler step to update the state
        variables. This suffices as a proof of concept.

        :param float dt: time step in seconds.
        :param int maxsteps: maximum number of sub-steps to simulate for one
                             ``dt`` update.
        """
        self.processCommandQueue()

        # Fetch the forces for all object positions.
        idPos = {k: v.position for (k, v) in self.allObjects.items()}
        ret = self.getGridForces(idPos)
        if not ret.ok:
            self.logit.info(ret.msg)
            z = np.float64(0)
            gridForces = {_: z for _ in idPos}
        else:
            gridForces = ret.data
        del ret, idPos

        # Iterate over all objects and update their SV information in Bullet.
        for objID, sv in self.allObjects.items():
            # Compute direct- and booster forces on object.
            force, torque = self.totalForceAndTorque(objID)

            # Add the force defined on the 'force' grid.
            force += gridForces[objID]

            # Update velocity and position.
            vel = np.array(sv.velocityLin, np.float64) + 0.5 * force
            pos = np.array(sv.position, np.float64)
            sv.velocityLin[:] = vel.tolist()
            sv.position[:] = (pos + dt * vel).tolist()

            self.allObjects[objID] = sv

        # Synchronise the local object cache back to the database.
        self.syncObjects(writeconcern=False)

    def processCommandQueue(self):
        """
        Apply commands from queue to objects in local cache.

        Applied commands are automatically removed.

        :return bool: Success.
        """
        # Fetch (and de-queue) all pending commands.
        ret = leoAPI.dequeueCommands()
        if not ret.ok:
            msg = 'Cannot fetch commands'
            self.logit.error(msg)
            return RetVal(False, msg, None)

        # Convenience.
        cmds = ret.data

        # Remove objects.
        for doc in cmds['remove']:
            objID = doc['objID']
            if objID in self.allObjects:
                self._DB_SV.remove({'objID': objID})
                del self.allObjects[objID]
                del self.allForces[objID]
                del self.allAABBs[objID]

        # Spawn objects.
        for doc in cmds['spawn']:
            objID = doc['objID']
            if objID in self.allObjects:
                msg = 'Cannot spawn object since objID={} already exists'
                self.logit.warning(msg.format(objID))
            else:
                sv_old = doc['sv']
                self.allObjects[objID] = _RigidBodyState(*sv_old)
                self.allForces[objID] = Forces(*(([0, 0, 0], ) * 4))
                self.allAABBs[objID] = float(doc['AABB'])

        # Update Body States.
        fun = leoAPI._updateRigidBodyStateTuple
        for doc in cmds['modify']:
            objID, sv_new = doc['objID'], doc['sv']
            if objID in self.allObjects:
                sv_new = RigidBodyStateOverride(*sv_new)
                sv_old = self.allObjects[objID]
                sv_old = RigidBodyState(*sv_old)
                self.allObjects[objID] = fun(sv_old, sv_new)

        # Update direct force- and torque values.
        for doc in cmds['direct_force']:
            objID, force, torque = doc['objID'], doc['force'], doc['torque']
            try:
                self.allForces[objID] = self.allForces[objID]._replace(
                    forceDirect=force, torqueDirect=torque)
            except KeyError:
                pass

        # Update booster force- and torque values.
        for doc in cmds['booster_force']:
            objID, force, torque = doc['objID'], doc['force'], doc['torque']
            try:
                self.allForces[objID] = self.allForces[objID]._replace(
                    forceBoost=force, torqueBoost=torque)
            except KeyError:
                pass

        return RetVal(True, None, None)

    def syncObjects(self, writeconcern: bool):
        """
        Sync the local BodyStates to the DB.

        The ``writeconcern`` flag is mostly for performance tuning. If set to
        *False* then the sync will not wait for an acknowledgement from the
        database after the write opration.

        :param bool writeconcern: disable write concern when set to *False*.
        """
        # Return immediately if we have no objects to begin with.
        if len(self.allObjects) == 0:
            return

        # Update (or insert non-existing) bodies. Use a MongoDB Bulk operator
        # for the update to improve the performance.
        bulk = self._DB_SV.initialize_unordered_bulk_op()
        for objID, sv in self.allObjects.items():
            query = {'objID': objID}
            data = {'objID': objID, 'sv': sv, 'AABB': self.allAABBs[objID]}
            bulk.find(query).upsert().update({'$set': data})

        if writeconcern:
            bulk.execute()
        else:
            bulk.execute({'w': 0, 'j': False})

    def processCommandsAndSync(self):
        """
        Process all pending commands and syncronise the cache to the DB.

        This method is useful for unit tests but probably not much else. It
        also ensures that the synchronisation of the objects from the local
        cache to the database is acknowledged by the database (this is turned
        off in production because it slows down Leonard an the information is
        not critical).
        """
        self.processCommandQueue()
        self.syncObjects(writeconcern=True)

    def run(self):
        """
        Drive the periodic physics updates.
        """
        # Call `run` method of `AzraelProcess` base class.
        super().run()

        # Initialisation.
        self.setup()
        self.logit.debug('Setup complete.')

        # Trigger the `step` method every `stepinterval` seconds, if possible.
        t0 = time.time()
        stepinterval = 0.050

        while True:
            # Wait, if less than 10ms have passed, or proceed immediately.
            sleep_time = stepinterval - (time.time() - t0)
            if sleep_time > 0:
                time.sleep(sleep_time)

            # Backup the time stamp.
            t0 = time.time()

            # Trigger the physics update step.
            with util.Timeit('Leonard:1.0 Step') as timeit:
                self.step(stepinterval, 10)


class LeonardBullet(LeonardBase):
    """
    An extension of ``LeonardBase`` that uses Bullet for the physics.

    Unlike ``LeonardBase`` this class actually *does* update the physics.
    """
    def __init__(self):
        super().__init__()
        self.bullet = None

    def setup(self):
        # Instantiate the Bullet engine with ID=1.
        self.bullet = azrael.bullet_api.PyBulletDynamicsWorld(1)

    @typecheck
    def step(self, dt, maxsteps):
        """
        Advance the simulation by ``dt`` using at most ``maxsteps``.

        This method will query all SV objects from the database and updates
        them in the Bullet engine. Then it defers to Bullet for the physics
        update.  Finally it copies the updated values in Bullet back to the
        database.

        :param float dt: time step in seconds.
        :param int maxsteps: maximum number of sub-steps to simulate for one
                             ``dt`` update.
        """
        # Convenience.
        vg = azrael.vectorgrid

        # Process pending commands.
        self.processCommandQueue()

        # Update the constraint cache in our local Igor instance.
        self.igor.updateLocalCache()
        allConstraints = self.igor.getAllConstraints().data

        # Fetch the forces for all object positions.
        idPos = {k: v.position for (k, v) in self.allObjects.items()}
        ret = self.getGridForces(idPos)
        if not ret.ok:
            self.logit.info(ret.msg)
            z = np.float64(0)
            gridForces = {_: z for _ in idPos}
        else:
            gridForces = ret.data
        del ret, idPos

        # Iterate over all objects and update them.
        for objID, sv in self.allObjects.items():
            # Pass the SV data from the DB to Bullet.
            self.bullet.setRigidBodyData(objID, sv)

            # Compute direct- and booster forces on object.
            force, torque = self.totalForceAndTorque(objID)

            # Add the force defined on the 'force' grid.
            force += gridForces[objID]

            # Apply the force to the object.
            self.bullet.applyForceAndTorque(objID, force, torque)

        # Apply all constraints. Log any errors but ignore them otherwise as
        # they are harmless (simply means no constraints were applied).
        ret = self.bullet.setConstraints(allConstraints)
        if not ret.ok:
            self.logit.warning(ret.msg)

        # Advance the simulation by one time step.
        with util.Timeit('compute') as timeit:
            self.bullet.compute(list(self.allObjects.keys()), dt, maxsteps)

            # Remove all constraints.
            self.bullet.clearAllConstraints()

        # Retrieve all objects from Bullet and overwrite the state variables
        # that the user explicilty wanted to change (if any).
        for objID in self.allObjects:
            ret = self.bullet.getRigidBodyData(objID)
            if ret.ok:
                self.allObjects[objID] = ret.data

        # Synchronise the local object cache back to the database.
        self.syncObjects(writeconcern=False)


class LeonardSweeping(LeonardBase):
    """
    Compute physics on independent collision sets.

    This is a modified version of ``LeonardBullet`` that uses
    Sweeping to compile the collision sets and then updates the physics for
    each set independently.

    This class is single threaded and uses a single Bullet instance to
    sequentially update the physics for each collision set.
    """
    def setup(self):
        # Instantiate the Bullet engine with ID=1.
        self.bullet = azrael.bullet_api.PyBulletDynamicsWorld(1)

    @typecheck
    def step(self, dt, maxsteps):
        """
        Advance the simulation by ``dt`` using at most ``maxsteps``.

        This method will query all SV objects from the database and updates
        them in the Bullet engine. Then it defers to Bullet for the physics
        update.  Finally it copies the updated values in Bullet back to the
        database.

        :param float dt: time step in seconds.
        :param int maxsteps: maximum number of sub-steps to simulate for one
                             ``dt`` update.
        """
        self.processCommandQueue()

        # Update the constraint cache in our local Igor instance.
        self.igor.updateLocalCache()
        allConstraints = self.igor.getAllConstraints().data

        # Compute all collision sets.
        with util.Timeit('CCS') as timeit:
            ret = self.igor.uniquePairs()
            if not ret.ok:
                return
            uniquePairs = ret.data

            ret = getFinalCollisionSets(
                uniquePairs, self.allObjects, self.allAABBs)
            if not ret.ok:
                return
            collSets = ret.data
            del ret, uniquePairs

        # Log the number of created collision sets.
        util.logMetricQty('#CollSets', len(collSets))

        # Convenience.
        vg = azrael.vectorgrid

        # Process all subsets individually.
        for subset in collSets:
            # Compile the subset dictionary for the current collision set.
            coll_SV = {_: self.allObjects[_] for _ in subset}

            # Fetch the forces for all object positions.
            idPos = {k: v.position for (k, v) in coll_SV.items()}
            ret = self.getGridForces(idPos)
            if not ret.ok:
                self.logit.info(ret.msg)
                z = np.float64(0)
                gridForces = {_: z for _ in idPos}
            else:
                gridForces = ret.data
            del ret, idPos

            # Iterate over all objects and update them.
            constr = []
            for objID, sv in coll_SV.items():
                # Pass the SV data from the DB to Bullet.
                self.bullet.setRigidBodyData(objID, sv)

                # Compute direct- and booster forces on object.
                force, torque = self.totalForceAndTorque(objID)

                # Add the force defined on the 'force' grid.
                force += gridForces[objID]

                # Apply the final force to the object.
                self.bullet.applyForceAndTorque(objID, force, torque)

            # Query all constraints and apply them in the next step (this
            # duplicates the code from LeonardBullet but I do not know of a
            # simple way to avoid it without major changes to the class
            # structure - for now this is acceptable, especially because this
            # class is mostly for testing).
            tmp = self.igor.getConstraints(coll_SV.keys()).data

            # Apply all constraints. Log any errors but ignore them otherwise
            # as they are harmless (simply means no constraints were applied).
            ret = self.bullet.setConstraints(tmp)
            if not ret.ok:
                self.logit.warning(ret.msg)
            del tmp

            # Wait for Bullet to advance the simulation by one step.
            with util.Timeit('compute') as timeit:
                self.bullet.compute(list(coll_SV.keys()), dt, maxsteps)

            # Remove all constraints.
            self.bullet.clearAllConstraints()

            # Retrieve all objects from Bullet.
            for objID, sv in coll_SV.items():
                ret = self.bullet.getRigidBodyData(objID)
                if ret.ok:
                    self.allObjects[objID] = ret.data

        # Synchronise the local object cache back to the database.
        self.syncObjects(writeconcern=False)


class LeonardDistributedZeroMQ(LeonardBase):
    """
    Compute physics with separate engines.

    This class uses the concept of Work Packages to distribute work. Every Work
    Package is self contained and holds all the information Bullet requires to
    step the simulation.

    This class uses the sweeping algorithm to determine collision sets, just
    like ``LeonardSweeping`` does.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.numWorkers = 3
        self.wpid_counter = 0
        self.ctx = None
        self.sock = None

        # Worker terminate automatically after a certain number of processed
        # Work Packages. The precise number is a constructor argument and the
        # following two variables simply specify the range. The final number
        # will be chosen randomly from this interval (different for every
        # Worker instance to minimise the probability that all terminate
        # simultaneously).
        self.minSteps, self.maxSteps = (500, 700)

        # Initialise worker manager handle.
        self.workermanager = None

    def __del__(self):
        """
        Kill all worker processes.
        """
        if self.workermanager is not None:
            # Send Keyboard interrupt to Worker manager and try to join it.
            os.kill(self.workermanager.pid, signal.SIGINT)
            self.workermanager.join(1)

            # Kill the Worker Manager if we could not join it.
            if self.workermanager.is_alive():
                self.workermanager.terminate()
                self.workermanager.join()

        # Close the Leonard <---> Worker socket.
        if self.sock is not None:
            self.sock.unbind(config.addr_leonard_repreq)
            self.sock.close(linger=100)

        # Terminate the ZeroMQ context.
        if self.ctx is not None:
            self.ctx.term()
            self.ctx.destroy()

    def setup(self):
        # Start the WorkerManager which, in turn, will spawn the workers.
        self.workermanager = WorkerManager(
            self.numWorkers, self.minSteps,
            self.maxSteps, LeonardWorkerZeroMQ)
        self.workermanager.start()

        # Create the ZeroMQ context. Note: this MUST happen AFTER the Worker
        # Manager was started, because otherwise the child processes will get a
        # copy of the ZeroMQ context with the already bound address, which it
        # may never release.
        self.ctx = zmq.Context()
        self.sock = self.ctx.socket(zmq.REP)

        # Bind the socket to the specified address. Retry a few times if
        # necessary.
        for ii in range(10):
            try:
                self.sock.bind(config.addr_leonard_repreq)
                break
            except zmq.error.ZMQError:
                time.sleep(0.2)
            assert ii < 9
        self.logit.info('Setup complete')

    @typecheck
    def step(self, dt, maxsteps):
        """
        Advance the simulation by ``dt`` using at most ``maxsteps``.

        This method moves all SV objects from the database to the Bullet
        engine. Then it defers to Bullet for the physics update. Finally, it
        replaces the SV fields with the user specified values (only applies if
        the user called 'setBodyState') and writes the results back to the
        database.

        :param float dt: time step in seconds.
        :param int maxsteps: maximum number of sub-steps to simulate for one
                             ``dt`` update.
        """
        # Read queued commands and update the local object cache accordingly.
        with util.Timeit('Leonard:1.1  processCmdQueue') as timeit:
            self.processCommandQueue()

        # Update the constraint cache in our local Igor instance.
        self.igor.updateLocalCache()
        allConstraints = self.igor.getAllConstraints().data

        # Compute the collision sets.
        with util.Timeit('Leonard:1.2  CCS') as timeit:
            ret = self.igor.uniquePairs()
            if not ret.ok:
                return
            uniquePairs = ret.data

            ret = getFinalCollisionSets(
                uniquePairs, self.allObjects, self.allAABBs)
            if not ret.ok:
                return
            collSets = ret.data
            del ret, uniquePairs

        # Log the number of created collision sets.
        util.logMetricQty('#CollSets', len(collSets))

        # Put each collision set into its own Work Package.
        with util.Timeit('Leonard:1.3  CreateWPs') as timeit:
            all_WPs = {}
            for subset in collSets:
                # Compile the Work Package. Skip this physics step altogether
                # if an error occurs.
                ret = self.createWorkPackage(list(subset), dt, maxsteps)
                if not ret.ok:
                    self.logit.error(ret.msg)
                    return
                all_WPs[ret.data['wpid']] = ret.data

        with util.Timeit('Leonard:1.4  WPSendRecv') as timeit:
            wpIdx = 0
            worklist = list(all_WPs.keys())
            while True:
                # Wait until Worker contacts us. The message usually contains a
                # processed Work Package. However, it may also be empty, most
                # likely because the Worker has not received a Work Package
                # from us yet.
                msg = self.sock.recv()
                if msg != b'':
                    # Unpickle the message.
                    msg = pickle.loads(msg)
                    wpid = msg['wpid']

                    # Ignore the message if its Work Package is not pending
                    # anymore (most likely because multiple Workers processed
                    # the same Work Package and the other Worker has already
                    # returned it).
                    if wpid in all_WPs:
                        self.updateLocalCache(msg['wpdata'])

                        # Decrement the Work Package Index if the wpIdx counter
                        # is already past that work package. This simply
                        # ensures that no WP is skipped simply because the
                        # queue has shrunk.
                        if worklist.index(wpid) < wpIdx:
                            wpIdx -= 1

                        # Remove the WP from the work list and the WP cache.
                        worklist.remove(wpid)
                        del all_WPs[wpid]

                # Send an empty message to the Worker if no Work Packages are
                # pending anymore. This empty message is important to avoid
                # breaking the REQ/REP pattern enforced by ZeroMQ.
                if len(all_WPs) == 0:
                    self.sock.send(b'')
                    break

                # Pick the next pending Work Package and increment the index.
                if wpIdx >= len(worklist):
                    wpIdx = 0
                wp = all_WPs[worklist[wpIdx]]
                wpIdx += 1

                # Send the Work Package to the Worker.
                self.sock.send(pickle.dumps(wp))

        # Synchronise the local cache back to the database.
        with util.Timeit('Leonard:1.5  syncObjects') as timeit:
            self.syncObjects(writeconcern=False)

    @typecheck
    def createWorkPackage(self, objIDs: (tuple, list),
                          dt: (int, float), maxsteps: int):
        """
        Create a new Work Package (WP) and return its ID.

        The Work Package will not be returned but uploaded to the DB directly.

        A Work Package carries the necessary information for another rigid body
        physics steps. The Worker can thus start its work immediately, at least
        in terms of the rigid bodies (it may still want to incorporate other
        data like grid forces, but that is up to the Worker implementation).

        The ``dt`` and ``maxsteps`` arguments are for the underlying physics
        engine.

        .. note::
           A work package contains only the objIDs but not their SV. The
           ``getNextWorkPackage`` function takes care of compiling this
           information.

        :param iterable objIDs: list of object IDs in the new work package.
        :param float dt: time step for this work package.
        :param int maxsteps: number of sub-steps for the time step.
        :return: Work package ID
        :rtype: int
        """
        # Sanity check.
        if len(objIDs) == 0:
            return RetVal(False, 'Work package is empty', None)

        # Compile the Body States and forces into a list of ``WPData`` tuples.
        try:
            wpdata = []
            for objID in objIDs:
                sv = self.allObjects[objID]
                force, torque = self.totalForceAndTorque(objID)
                wpdata.append(WPData(objID, sv, force, torque))
        except KeyError as err:
            return RetVal(False, 'Cannot compile WP', None)

        # Query all constraints.
        constraints = self.igor.getConstraints(objIDs).data

        # Form the content of the Work Package as it will appear in the DB.
        data = {'wpid': self.wpid_counter,
                'wpmeta': (self.wpid_counter, dt, maxsteps),
                'wpdata': wpdata,
                'wpconstraints': constraints,
                'ts': None}
        self.wpid_counter += 1
        return RetVal(True, None, data)

    def updateLocalCache(self, wpdata):
        """
        Copy every object from ``wpdata`` to the local cache.

        The ``wpdata`` argument is a list of (objID, sv) tuples.

        The implicit assumption of this method is that ``wpdata`` is the
        output of ``computePhysicsForWorkPackage`` from a Worker.

        :param list wpdata: Content of Work Packge as returned by Workers.
        """
        # Reset force and torque for all objects in the WP, and overwrite
        # the old Body States with the new one from the processed WP.
        for (objID, sv) in wpdata:
            self.allObjects[objID] = _RigidBodyState(*sv)


class LeonardWorkerZeroMQ(config.AzraelProcess):
    """
    Dedicated Worker to process Work Packages.

    :param int workerID: the ID of this worker.
    :param int stepsUntilQuit: Worker will restart after this many steps.
    """
    def __init__(self, workerID, stepsUntilQuit: int):
        super().__init__()
        self.workerID = workerID

        # After ``stepsUntilQuit`` this Worker will spawn a new Worker with the
        # same ID and quit.
        assert stepsUntilQuit > 0
        self.stepsUntilQuit = stepsUntilQuit

        # Instantiate a Bullet engine.
        engine = azrael.bullet_api.PyBulletDynamicsWorld
        self.bullet = engine(self.workerID)

    def getGridForces(self, idPos: dict):
        """
        Return dictionary of force values for every object in ``idPos``.

        The ``idPos`` argument is a {objID_1: sv_1, objID_2, sv_2, ...}
        dictionary.

        The returned dictionary has the same keys as ``idPos``.

        :param dict idPos: dictionary with objIDs and corresponding SVs.
        :return dict: {objID_k: force_k}
        """
        # Convenience.
        vg = azrael.vectorgrid

        # Extract the keys and values in the same order.
        objIDs = list(idPos.keys())
        positions = [idPos[_] for _ in objIDs]

        # Query the grid values at all positions.
        ret = vg.getValues('force', positions)
        if not ret.ok:
            return RetVal(False, ret.msg, None)

        # Overwrite the default values with whatever ...
        gridForces = {objID: val for objID, val in zip(objIDs, ret.data)}
        return RetVal(True, None, gridForces)

    def computePhysicsForWorkPackage(self, wp):
        """
        Compute a physics steps for all objects in ``wp``.

        The output of this method is matched to the ``updateLocalCache`` in
        Leonard itself.

        :param dict wp: Work Package content from ``createWorkPackage``.
        :return dict: {'wpdata': list_of_SVs, 'wpid': wpid}
        """
        worklist, meta = wp['wpdata'], WPMeta(*wp['wpmeta'])
        constraints = wp['wpconstraints']

        # Log the number of collision-sets in the current Work Package.
        util.logMetricQty('Engine_{}'.format(self.workerID), len(worklist))

        # Convenience: the WPData elements make the code more readable.
        worklist = [WPData(*_) for _ in worklist]

        # Convenience.
        applyForceAndTorque = self.bullet.applyForceAndTorque
        setRigidBodyData = self.bullet.setRigidBodyData

        # Add every object to the Bullet engine and set the force/torque.
        with util.Timeit('Worker:1.1.0  applyforce') as timeit:
            with util.Timeit('Worker:1.1.1   grid') as timeit:
                # Fetch the forces for all object positions.
                idPos = {_.aid: _.sv.position for _ in worklist}
                ret = self.getGridForces(idPos)
                if not ret.ok:
                    self.logit.info(ret.msg)
                    z = np.float64(0)
                    gridForces = {_: z for _ in idPos}
                else:
                    gridForces = ret.data
                del ret, idPos

            with util.Timeit('Worker:1.1.1   updateGeo') as timeit:
                for obj in worklist:
                    # Update the object in Bullet and apply the force/torque.
                    setRigidBodyData(obj.aid, obj.sv)

            with util.Timeit('Worker:1.1.1   updateForce') as timeit:
                for obj in worklist:
                    # Add the force defined on the 'force' grid.
                    force = obj.force + gridForces[obj.aid]
                    applyForceAndTorque(obj.aid, force, obj.torque)

        # Apply all constraints. Log any errors but ignore them otherwise as
        # they are harmless (simply means no constraints were applied).
        ret = self.bullet.setConstraints(constraints)
        if not ret.ok:
            self.logit.warning(ret.msg)

        # Tell Bullet to advance the simulation for all objects in the
        # current work list.
        with util.Timeit('Worker:1.2.0  compute') as timeit:
            IDs = [_.aid for _ in worklist]
            self.bullet.compute(IDs, meta.dt, meta.maxsteps)

            # Remove all constraints.
            self.bullet.clearAllConstraints()

        with util.Timeit('Worker:1.3.0  fetchFromBullet') as timeit:
            # Retrieve the objects from Bullet again and update them in the DB.
            out = []
            for obj in worklist:
                ret = self.bullet.getRigidBodyData(obj.aid)
                sv = ret.data
                if not ret.ok:
                    # Something went wrong. Reuse the old SV.
                    sv = obj.sv
                    self.logit.error('Unable to get all objects from Bullet')
                out.append((obj.aid, sv))

        # Return the updated WP data.
        return {'wpid': meta.wpid, 'wpdata': out}

    @typecheck
    def run(self):
        """
        Wait for Work Packages, process them, and return the results.
        """
        # Call `run` method of `AzraelProcess` base class.
        super().run()

        try:
            # Setup ZeroMQ.
            ctx = zmq.Context()
            sock = ctx.socket(zmq.REQ)
            sock.connect(config.addr_leonard_repreq)
            self.logit.info('Worker {} connected'.format(self.workerID))

            # Contact Leonard with an empty payload.
            sock.send(b'')

            # Wait for messages from Leonard. If they contain a WP then process
            # it and return the result, otherwise reply with an empty message.
            numSteps = 0
            suq = self.stepsUntilQuit
            while numSteps < suq:
                # Wait for the next message.
                msg = sock.recv()

                # If Leonard did not send a Work Package (probably because it
                # does not have one right now) then wait for a short time
                # before asking again to avoid spamming the network.
                if msg == b'':
                    time.sleep(0.003)
                    sock.send(b'')
                    continue

                # Unpickle the Work Package.
                wpdata = pickle.loads(msg)

                # Process the Work Package.
                with util.Timeit('Worker:1.0.0 WPTotal') as timeit:
                    wpdata = self.computePhysicsForWorkPackage(wpdata)

                # Pack up the Work Package and send it back to Leonard.
                sock.send(pickle.dumps(wpdata))

                # Count the number of Work Packages we have processed.
                numSteps += 1

            # Log a last status message before terminating.
            self.logit.info('Worker {} terminated itself after {} steps'
                            .format(self.workerID, numSteps))
        except KeyboardInterrupt:
            print('Aborted Worker {}'.format(self.workerID))

        # Terminate.
        sock.close(linger=0)
        ctx.destroy()
        print('Worker {} exited cleanly'.format(self.workerID))


class WorkerManager(config.AzraelProcess):
    """
    Launch Worker processes and restart them as necessary.

    This class merely launches the inital set of workers and periodically
    checks if any have died. If so, it joins these processes and replaces it
    with a new Worker that has the same ID.

    :param int numWorker: number of Workers processes to spawn.
    :param int minSteps: see Worker
    :param int maxSteps: see Worker
    :param class workerCls: the class to instantiate.
    """
    def __init__(self, numWorkers: int, minSteps: int, maxSteps: int,
                 workerCls):
        super().__init__()

        # Sanity checks.
        assert numWorkers > 0
        assert isinstance(minSteps, int)
        assert isinstance(maxSteps, int)
        assert 0 < minSteps <= maxSteps

        # Backup the arguments.
        self.numWorkers = numWorkers
        self.workerCls = workerCls
        self.minSteps, self.maxSteps = minSteps, maxSteps

        # Initialise the list of workers.
        self.workers = []

    def _run(self):
        """
        Start the initial collection of Workers and ensure they remain alive.
        """
        # Spawn the initial collection of Workers.
        delta = self.maxSteps - self.minSteps
        for ii in range(self.numWorkers):
            # Random number in [minSteps, maxSteps]. The process will
            # automatically terminate after `suq` steps.
            suq = self.minSteps + int(np.random.rand() * delta)

            # Instantiate the process and add it to the list.
            self.workers.append(self.workerCls(ii + 1, suq))
            self.workers[-1].start()

        # Periodically monitor the processes and restart any that have died.
        while True:
            # Only check once a second.
            time.sleep(1)
            for workerID, proc in enumerate(self.workers):
                # Skip current process if it is still running.
                if proc.is_alive():
                    continue

                # Process has died --> join it to clear up the process table.
                proc.join()

                # Create a new Worker with the same ID but a (possibly)
                # different number of steps after which it must terminate.
                suq = self.minSteps + int(np.random.rand() * delta)
                proc = self.workerCls(workerID, suq)
                proc.start()
                self.workers[workerID] = proc
                self.logit.info('Restarted Worker {}'.format(workerID))

    def run(self):
        """
        Wrapper around ``_run`` to intercept SIGTERM.
        """
        # Call `run` method of `AzraelProcess` base class.
        super().run()

        try:
            self._run()
        except KeyboardInterrupt:
            print('Worker Manager was aborted')

        # Terminate all workers.
        for w in self.workers:
            if w is not None and w.is_alive():
                w.terminate()
                w.join()
        print('Worker manager finished')
<commit_msg>'computeCollisionSetsAABB' now uses AABBs prisms instead of cubes.<commit_after># Copyright 2014, O<NAME> <<EMAIL>>
#
# This file is part of Azrael (https://github.com/olitheolix/azrael)
#
# Azrael is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# Azrael is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with Azrael. If not, see <http://www.gnu.org/licenses/>.

"""
Physics manager.
"""
import os
import sys
import zmq
import time
import signal
import pickle
import itertools
import numpy as np

import azrael.igor
import azrael.database
import azrael.vectorgrid
import azrael.bullet_api
import azrael.util as util
import azrael.config as config
import azrael.rb_state as rb_state
import azrael.leo_api as leoAPI

from IPython import embed as ipshell
from azrael.types import _RigidBodyState
from azrael.types import typecheck, RetVal, WPData, WPMeta, Forces

# Convenience.
RigidBodyState = rb_state.RigidBodyState
RigidBodyStateOverride = rb_state.RigidBodyStateOverride


@typecheck
def sweeping(data: list, labels: np.ndarray, dim: str):
    """
    Return sets of overlapping AABBs in the dimension ``dim``.

    This function implements the 'Sweeping' algorithm to determine which sets
    of AABBs overlap.

    Sweeping is straightforward: sort all start/stop positions and determine
    the overlapping sets.

    The returned sets do not contain the ``data`` elements but their
    corresponding ``labels`` to be more memory efficient.

    The `data` is a list of dictionaries that denotes the half widths of the
    AABBs in the respective direction:
      data = [{'x': [x0, x1], 'y': [y0, y1], 'z': [z0, z1]},
              {'x': [x0, x1], 'y': [y0, y1], 'z': [z0, z1]},
              ...]
    The elements denote the AABB of the corresponding `label` at the same
    index. The dictionary keys 'x', 'y', and 'z' are (more or less) hard coded,
    and `dim` must refer to one of them (eg dim = 'y' to find the sets that
    overlap in the 'y' dimension).

    :param list[dict] data: list of dictionaries containing the AABB boundaries
       for 'x', 'y', and 'z'.
    :param np.int64 labels: integer array to label the elements in data.
    :param str dim: the axis to check (must be one of ['x', 'y', 'z'])
    :return: list of sets. Each set contains elements from ``labels``.
    :rtype: list of sets
    """
    assert len(labels) == len(data)

    # Convenience.
    N = 2 * len(data)

    # Pre-allocate arrays for start/stop position, objID, and an
    # increment/decrement array used for convenient processing afterwards.
    arr_pos = np.zeros(N, np.float64)
    arr_lab = np.zeros(N, np.int64)
    arr_inc = np.zeros(N, np.int8)

    # Fill the arrays.
    for ii in range(len(data)):
        arr_pos[2 * ii: 2 * ii + 2] = np.array(data[ii][dim])
        arr_lab[2 * ii: 2 * ii + 2] = labels[ii]
        arr_inc[2 * ii: 2 * ii + 2] = [+1, -1]

    # Sort all three arrays according to the start/stop positions.
    idx = np.argsort(arr_pos)
    arr_lab = arr_lab[idx]
    arr_inc = arr_inc[idx]

    # Output array.
    out = []

    # Sweep over the sorted data and compile the list of object sets.
    sumVal = 0
    setObjs = set()
    for (inc, objID) in zip(arr_inc, arr_lab):
        # Update the index variable and add the current object to the set.
        sumVal += inc
        setObjs.add(objID)

        # A new set of overlapping AABBs is complete whenever `sumVal`
        # reaches zero.
        if sumVal == 0:
            out.append(setObjs)
            setObjs = set()

        # Safety check: sumVal can never be negative.
        assert sumVal >= 0
    return RetVal(True, None, out)


@typecheck
def computeCollisionSetsAABB(SVs: dict, AABBs: dict):
    """
    Return broadphase collision sets for the objects in ``SVs``.

    :param dict SVs: Dictionary of Body States.
    :param dict AABBs: Dictionary of AABBs.
    :return: each list contains a unique set of overlapping objects.
    :rtype: list of lists
    """
    # Sanity check.
    if set(SVs.keys()) != set(AABBs.keys()):
        return RetVal(False, 'SVs and AABBs are inconsisten', None)

    # The 'sweeping' function requires a list of dictionaries. Each dictionary
    # must contain the min/max spatial extent in x/y/z direction.
    data = []
    IDs = list(SVs.keys())
    for objID in IDs:
        if (SVs[objID] is None) or (AABBs[objID] is None):
            continue
        sv, aabb = SVs[objID], AABBs[objID]
        pos = sv.position
        x0, x1 = pos[0] - aabb[0], pos[0] + aabb[0]
        y0, y1 = pos[1] - aabb[1], pos[1] + aabb[1]
        z0, z1 = pos[2] - aabb[2], pos[2] + aabb[2]

        data.append({'x': [x0, x1], 'y': [y0, y1], 'z': [z0, z1]})
    del SVs, AABBs

    # Enumerate the objects.
    labels = np.arange(len(IDs))

    # Determine the sets of objects that overlap 'x' direction.
    stage_0 = sweeping(data, labels, 'x').data

    # Iterate over all the just found sets. For each, determine the sets that
    # overlap in the 'y' dimension.
    stage_1 = []
    for subset in stage_0:
        tmpData = [data[_] for _ in subset]
        tmpLabels = np.array(tuple(subset), np.int64)
        stage_1.extend(sweeping(tmpData, tmpLabels, 'y').data)

    # Iterate over all the sets that overlap in 'x' and 'y' dimension. For
    # each, determine which also overalp in the 'z' dimension.
    stage_2 = []
    for subset in stage_1:
        tmpData = [data[_] for _ in subset]
        tmpLabels = np.array(tuple(subset), np.int64)
        stage_2.extend(sweeping(tmpData, tmpLabels, 'z').data)

    # Convert the labels back to object IDs.
    out = [[IDs[objID] for objID in objIDs] for objIDs in stage_2]

    # Return the sets.
    return RetVal(True, None, out)


def mergeConstraintSets(constraintPairs: tuple, collSets: (tuple, list)):
    """
    Merge all the sets in ``collSets`` that contain any of the
    ``constraintPairs``.

    The purpose of this function is to merge those collision sets that are
    connected via a constraint. Typically, this function takes the output of
    ``computeCollisionSets`` as the ``collSets`` argument.

    :param list constraintPairs: list of 2-tuples eg [(1, 2), (1, 5), ...].
    :param list collSets: a set/list of sets/lists.
    :return: list of collision sets.
    """
    # Merge the collision sets that are linked via constraints.
    for (a, b) in constraintPairs:
        # Find (and remove) the set(s) that contain object_a or object_b (or
        # both).
        s_a = [collSets.pop(ii) for ii, v in enumerate(collSets) if a in v]
        s_b = [collSets.pop(ii) for ii, v in enumerate(collSets) if b in v]

        # Sanity check: each element can be in at most one set.
        assert 0 <= len(s_a) < 2
        assert 0 <= len(s_b) < 2

        # Remove the list from the list comprehension.
        s_a = s_a[0] if len(s_a) == 1 else []
        s_b = s_b[0] if len(s_b) == 1 else []

        # Merge the collision sets.
        collSets.append(s_a + s_b)
    return RetVal(True, None, collSets)


def getFinalCollisionSets(
        constraintPairs: list, allObjects: list, allAABBs: list):
    """
    Return the collision sets.

    This function calls the broadphase solver and then merges those collision
    sets that are connected via a constraint.

    .. note:: this function is tightly coupled to the inner working of Leonard.
              Possibly it should not be a standalone function.

    :param list constraintPairs: list of 2-tuples eg [(1, 2), (1, 5), ...].
    :param list allObjects: Leonard's object cache.
    :param list allAABBs: Leonard's AABB cache.
    :return: list of non-overlapping collision sets.
    """
    # Broadphase based on AABB only.
    ret = computeCollisionSetsAABB(allObjects, allAABBs)
    if not ret.ok:
        msg = 'ComputeCollisionSetsAABB returned an error'
        self.logit.error(msg)
        return RetVal(False, msg, None)

    # Merge all collision sets that have objects which are connected by a
    # constraint.
    collSets = ret.data
    ret = mergeConstraintSets(constraintPairs, collSets)
    if not ret.ok:
        msg = 'mergeConstraintSets returned an error'
        self.logit.error(msg)
        return RetVal(False, msg, None)
    return ret


class LeonardBase(config.AzraelProcess):
    """
    Base class for Physics manager.

    No physics is actually computed here. The class serves mostly as an
    interface for the actual Leonard implementations, as well as a test
    framework.
    """
    def __init__(self):
        super().__init__()

        # Create the DB handles.
        self._DB_SV = azrael.database.dbHandles['SV']

        # Create an Igor instance.
        self.igor = azrael.igor.Igor()

        self.allObjects = {}
        self.allAABBs = {}
        self.allForces = {}

    def setup(self):
        """
        Stub for initialisation code that cannot go into the constructor.

        This is typically necessary for code that must not execute until
        *after* the class forked into a new process. ZeroMQ contexts are a good
        example for this - *never* fork processes *after* creating the context.
        """
        pass

    def getGridForces(self, idPos: dict):
        """
        Return dictionary of force values for every object in ``idPos``.

        The ``idPos`` argument is a {objID_1: sv_1, objID_2, sv_2, ...}
        dictionary.

        The returned dictionary has the same keys as ``idPos``.

        :param dict idPos: dictionary with objIDs and corresponding SVs.
        :return dict: {objID_k: force_k}
        """
        # Convenience.
        vg = azrael.vectorgrid

        # Extract the keys and values in the same order.
        objIDs = list(idPos.keys())
        positions = [idPos[_] for _ in objIDs]

        # Query the grid values at all positions.
        ret = vg.getValues('force', positions)
        if not ret.ok:
            return RetVal(False, ret.msg, None)

        # Overwrite the default values.
        gridForces = {objID: val for objID, val in zip(objIDs, ret.data)}
        return RetVal(True, None, gridForces)

    @typecheck
    def totalForceAndTorque(self, objID: int):
        """
        Return the total force- and torque on the object.

        The returned values are the sum of all booster forces (correctly
        oriented relative to the object) and the force a user may have specifed
        directly.

        Note that this function does not account for the forces from the 'force
        grid'.

        :param int objID: return the force and torque for this object.
        :return: the force and torque as two Python lists (not NumPy arrays)
        :rtype: (list, list)
        """
        # Convenience.
        sv = self.allObjects[objID]
        f = self.allForces[objID]

        # Construct the Quaternion of the object based on its orientation.
        quat = util.Quaternion(sv.orientation[3], sv.orientation[:3])

        # Fetch the force vector for the current object from the DB.
        force = np.array(f.forceDirect, np.float64)
        torque = np.array(f.torqueDirect, np.float64)

        # Add the booster's contribution to force and torque.
        # Note: We cannot do this directly since the booster force/torque were
        # specified in object coordinates. We thus rotate them to world
        # coordinates before adding them to the total force.
        force += quat * f.forceBoost
        torque += quat * f.torqueBoost

        # Convert to Python lists.
        return force.tolist(), torque.tolist()

    @typecheck
    def step(self, dt: (int, float), maxsteps: int):
        """
        Advance the simulation by ``dt`` using at most ``maxsteps``.

        This method will use a primitive Euler step to update the state
        variables. This suffices as a proof of concept.

        :param float dt: time step in seconds.
        :param int maxsteps: maximum number of sub-steps to simulate for one
                             ``dt`` update.
        """
        self.processCommandQueue()

        # Fetch the forces for all object positions.
        idPos = {k: v.position for (k, v) in self.allObjects.items()}
        ret = self.getGridForces(idPos)
        if not ret.ok:
            self.logit.info(ret.msg)
            z = np.float64(0)
            gridForces = {_: z for _ in idPos}
        else:
            gridForces = ret.data
        del ret, idPos

        # Iterate over all objects and update their SV information in Bullet.
        for objID, sv in self.allObjects.items():
            # Compute direct- and booster forces on object.
            force, torque = self.totalForceAndTorque(objID)

            # Add the force defined on the 'force' grid.
            force += gridForces[objID]

            # Update velocity and position.
            vel = np.array(sv.velocityLin, np.float64) + 0.5 * force
            pos = np.array(sv.position, np.float64)
            sv.velocityLin[:] = vel.tolist()
            sv.position[:] = (pos + dt * vel).tolist()

            self.allObjects[objID] = sv

        # Synchronise the local object cache back to the database.
        self.syncObjects(writeconcern=False)

    def processCommandQueue(self):
        """
        Apply commands from queue to objects in local cache.

        Applied commands are automatically removed.

        :return bool: Success.
        """
        # Fetch (and de-queue) all pending commands.
        ret = leoAPI.dequeueCommands()
        if not ret.ok:
            msg = 'Cannot fetch commands'
            self.logit.error(msg)
            return RetVal(False, msg, None)

        # Convenience.
        cmds = ret.data

        # Remove objects.
        for doc in cmds['remove']:
            objID = doc['objID']
            if objID in self.allObjects:
                self._DB_SV.remove({'objID': objID})
                del self.allObjects[objID]
                del self.allForces[objID]
                del self.allAABBs[objID]

        # Spawn objects.
        for doc in cmds['spawn']:
            objID = doc['objID']
            if objID in self.allObjects:
                msg = 'Cannot spawn object since objID={} already exists'
                self.logit.warning(msg.format(objID))
            else:
                sv_old = doc['sv']
                self.allObjects[objID] = _RigidBodyState(*sv_old)
                self.allForces[objID] = Forces(*(([0, 0, 0], ) * 4))
                self.allAABBs[objID] = float(doc['AABB'])

        # Update Body States.
        fun = leoAPI._updateRigidBodyStateTuple
        for doc in cmds['modify']:
            objID, sv_new = doc['objID'], doc['sv']
            if objID in self.allObjects:
                sv_new = RigidBodyStateOverride(*sv_new)
                sv_old = self.allObjects[objID]
                sv_old = RigidBodyState(*sv_old)
                self.allObjects[objID] = fun(sv_old, sv_new)

        # Update direct force- and torque values.
        for doc in cmds['direct_force']:
            objID, force, torque = doc['objID'], doc['force'], doc['torque']
            try:
                self.allForces[objID] = self.allForces[objID]._replace(
                    forceDirect=force, torqueDirect=torque)
            except KeyError:
                pass

        # Update booster force- and torque values.
        for doc in cmds['booster_force']:
            objID, force, torque = doc['objID'], doc['force'], doc['torque']
            try:
                self.allForces[objID] = self.allForces[objID]._replace(
                    forceBoost=force, torqueBoost=torque)
            except KeyError:
                pass

        return RetVal(True, None, None)

    def syncObjects(self, writeconcern: bool):
        """
        Sync the local BodyStates to the DB.

        The ``writeconcern`` flag is mostly for performance tuning. If set to
        *False* then the sync will not wait for an acknowledgement from the
        database after the write opration.

        :param bool writeconcern: disable write concern when set to *False*.
        """
        # Return immediately if we have no objects to begin with.
        if len(self.allObjects) == 0:
            return

        # Update (or insert non-existing) bodies. Use a MongoDB Bulk operator
        # for the update to improve the performance.
        bulk = self._DB_SV.initialize_unordered_bulk_op()
        for objID, sv in self.allObjects.items():
            query = {'objID': objID}
            data = {'objID': objID, 'sv': sv, 'AABB': self.allAABBs[objID]}
            bulk.find(query).upsert().update({'$set': data})

        if writeconcern:
            bulk.execute()
        else:
            bulk.execute({'w': 0, 'j': False})

    def processCommandsAndSync(self):
        """
        Process all pending commands and syncronise the cache to the DB.

        This method is useful for unit tests but probably not much else. It
        also ensures that the synchronisation of the objects from the local
        cache to the database is acknowledged by the database (this is turned
        off in production because it slows down Leonard an the information is
        not critical).
        """
        self.processCommandQueue()
        self.syncObjects(writeconcern=True)

    def run(self):
        """
        Drive the periodic physics updates.
        """
        # Call `run` method of `AzraelProcess` base class.
        super().run()

        # Initialisation.
        self.setup()
        self.logit.debug('Setup complete.')

        # Trigger the `step` method every `stepinterval` seconds, if possible.
        t0 = time.time()
        stepinterval = 0.050

        while True:
            # Wait, if less than 10ms have passed, or proceed immediately.
            sleep_time = stepinterval - (time.time() - t0)
            if sleep_time > 0:
                time.sleep(sleep_time)

            # Backup the time stamp.
            t0 = time.time()

            # Trigger the physics update step.
            with util.Timeit('Leonard:1.0 Step') as timeit:
                self.step(stepinterval, 10)


class LeonardBullet(LeonardBase):
    """
    An extension of ``LeonardBase`` that uses Bullet for the physics.

    Unlike ``LeonardBase`` this class actually *does* update the physics.
    """
    def __init__(self):
        super().__init__()
        self.bullet = None

    def setup(self):
        # Instantiate the Bullet engine with ID=1.
        self.bullet = azrael.bullet_api.PyBulletDynamicsWorld(1)

    @typecheck
    def step(self, dt, maxsteps):
        """
        Advance the simulation by ``dt`` using at most ``maxsteps``.

        This method will query all SV objects from the database and updates
        them in the Bullet engine. Then it defers to Bullet for the physics
        update.  Finally it copies the updated values in Bullet back to the
        database.

        :param float dt: time step in seconds.
        :param int maxsteps: maximum number of sub-steps to simulate for one
                             ``dt`` update.
        """
        # Convenience.
        vg = azrael.vectorgrid

        # Process pending commands.
        self.processCommandQueue()

        # Update the constraint cache in our local Igor instance.
        self.igor.updateLocalCache()
        allConstraints = self.igor.getAllConstraints().data

        # Fetch the forces for all object positions.
        idPos = {k: v.position for (k, v) in self.allObjects.items()}
        ret = self.getGridForces(idPos)
        if not ret.ok:
            self.logit.info(ret.msg)
            z = np.float64(0)
            gridForces = {_: z for _ in idPos}
        else:
            gridForces = ret.data
        del ret, idPos

        # Iterate over all objects and update them.
        for objID, sv in self.allObjects.items():
            # Pass the SV data from the DB to Bullet.
            self.bullet.setRigidBodyData(objID, sv)

            # Compute direct- and booster forces on object.
            force, torque = self.totalForceAndTorque(objID)

            # Add the force defined on the 'force' grid.
            force += gridForces[objID]

            # Apply the force to the object.
            self.bullet.applyForceAndTorque(objID, force, torque)

        # Apply all constraints. Log any errors but ignore them otherwise as
        # they are harmless (simply means no constraints were applied).
        ret = self.bullet.setConstraints(allConstraints)
        if not ret.ok:
            self.logit.warning(ret.msg)

        # Advance the simulation by one time step.
        with util.Timeit('compute') as timeit:
            self.bullet.compute(list(self.allObjects.keys()), dt, maxsteps)

            # Remove all constraints.
            self.bullet.clearAllConstraints()

        # Retrieve all objects from Bullet and overwrite the state variables
        # that the user explicilty wanted to change (if any).
        for objID in self.allObjects:
            ret = self.bullet.getRigidBodyData(objID)
            if ret.ok:
                self.allObjects[objID] = ret.data

        # Synchronise the local object cache back to the database.
        self.syncObjects(writeconcern=False)


class LeonardSweeping(LeonardBase):
    """
    Compute physics on independent collision sets.

    This is a modified version of ``LeonardBullet`` that uses
    Sweeping to compile the collision sets and then updates the physics for
    each set independently.

    This class is single threaded and uses a single Bullet instance to
    sequentially update the physics for each collision set.
    """
    def setup(self):
        # Instantiate the Bullet engine with ID=1.
        self.bullet = azrael.bullet_api.PyBulletDynamicsWorld(1)

    @typecheck
    def step(self, dt, maxsteps):
        """
        Advance the simulation by ``dt`` using at most ``maxsteps``.

        This method will query all SV objects from the database and updates
        them in the Bullet engine. Then it defers to Bullet for the physics
        update.  Finally it copies the updated values in Bullet back to the
        database.

        :param float dt: time step in seconds.
        :param int maxsteps: maximum number of sub-steps to simulate for one
                             ``dt`` update.
        """
        self.processCommandQueue()

        # Update the constraint cache in our local Igor instance.
        self.igor.updateLocalCache()
        allConstraints = self.igor.getAllConstraints().data

        # Compute all collision sets.
        with util.Timeit('CCS') as timeit:
            ret = self.igor.uniquePairs()
            if not ret.ok:
                return
            uniquePairs = ret.data

            ret = getFinalCollisionSets(
                uniquePairs, self.allObjects, self.allAABBs)
            if not ret.ok:
                return
            collSets = ret.data
            del ret, uniquePairs

        # Log the number of created collision sets.
        util.logMetricQty('#CollSets', len(collSets))

        # Convenience.
        vg = azrael.vectorgrid

        # Process all subsets individually.
        for subset in collSets:
            # Compile the subset dictionary for the current collision set.
            coll_SV = {_: self.allObjects[_] for _ in subset}

            # Fetch the forces for all object positions.
            idPos = {k: v.position for (k, v) in coll_SV.items()}
            ret = self.getGridForces(idPos)
            if not ret.ok:
                self.logit.info(ret.msg)
                z = np.float64(0)
                gridForces = {_: z for _ in idPos}
            else:
                gridForces = ret.data
            del ret, idPos

            # Iterate over all objects and update them.
            constr = []
            for objID, sv in coll_SV.items():
                # Pass the SV data from the DB to Bullet.
                self.bullet.setRigidBodyData(objID, sv)

                # Compute direct- and booster forces on object.
                force, torque = self.totalForceAndTorque(objID)

                # Add the force defined on the 'force' grid.
                force += gridForces[objID]

                # Apply the final force to the object.
                self.bullet.applyForceAndTorque(objID, force, torque)

            # Query all constraints and apply them in the next step (this
            # duplicates the code from LeonardBullet but I do not know of a
            # simple way to avoid it without major changes to the class
            # structure - for now this is acceptable, especially because this
            # class is mostly for testing).
            tmp = self.igor.getConstraints(coll_SV.keys()).data

            # Apply all constraints. Log any errors but ignore them otherwise
            # as they are harmless (simply means no constraints were applied).
            ret = self.bullet.setConstraints(tmp)
            if not ret.ok:
                self.logit.warning(ret.msg)
            del tmp

            # Wait for Bullet to advance the simulation by one step.
            with util.Timeit('compute') as timeit:
                self.bullet.compute(list(coll_SV.keys()), dt, maxsteps)

            # Remove all constraints.
            self.bullet.clearAllConstraints()

            # Retrieve all objects from Bullet.
            for objID, sv in coll_SV.items():
                ret = self.bullet.getRigidBodyData(objID)
                if ret.ok:
                    self.allObjects[objID] = ret.data

        # Synchronise the local object cache back to the database.
        self.syncObjects(writeconcern=False)


class LeonardDistributedZeroMQ(LeonardBase):
    """
    Compute physics with separate engines.

    This class uses the concept of Work Packages to distribute work. Every Work
    Package is self contained and holds all the information Bullet requires to
    step the simulation.

    This class uses the sweeping algorithm to determine collision sets, just
    like ``LeonardSweeping`` does.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.numWorkers = 3
        self.wpid_counter = 0
        self.ctx = None
        self.sock = None

        # Worker terminate automatically after a certain number of processed
        # Work Packages. The precise number is a constructor argument and the
        # following two variables simply specify the range. The final number
        # will be chosen randomly from this interval (different for every
        # Worker instance to minimise the probability that all terminate
        # simultaneously).
        self.minSteps, self.maxSteps = (500, 700)

        # Initialise worker manager handle.
        self.workermanager = None

    def __del__(self):
        """
        Kill all worker processes.
        """
        if self.workermanager is not None:
            # Send Keyboard interrupt to Worker manager and try to join it.
            os.kill(self.workermanager.pid, signal.SIGINT)
            self.workermanager.join(1)

            # Kill the Worker Manager if we could not join it.
            if self.workermanager.is_alive():
                self.workermanager.terminate()
                self.workermanager.join()

        # Close the Leonard <---> Worker socket.
        if self.sock is not None:
            self.sock.unbind(config.addr_leonard_repreq)
            self.sock.close(linger=100)

        # Terminate the ZeroMQ context.
        if self.ctx is not None:
            self.ctx.term()
            self.ctx.destroy()

    def setup(self):
        # Start the WorkerManager which, in turn, will spawn the workers.
        self.workermanager = WorkerManager(
            self.numWorkers, self.minSteps,
            self.maxSteps, LeonardWorkerZeroMQ)
        self.workermanager.start()

        # Create the ZeroMQ context. Note: this MUST happen AFTER the Worker
        # Manager was started, because otherwise the child processes will get a
        # copy of the ZeroMQ context with the already bound address, which it
        # may never release.
        self.ctx = zmq.Context()
        self.sock = self.ctx.socket(zmq.REP)

        # Bind the socket to the specified address. Retry a few times if
        # necessary.
        for ii in range(10):
            try:
                self.sock.bind(config.addr_leonard_repreq)
                break
            except zmq.error.ZMQError:
                time.sleep(0.2)
            assert ii < 9
        self.logit.info('Setup complete')

    @typecheck
    def step(self, dt, maxsteps):
        """
        Advance the simulation by ``dt`` using at most ``maxsteps``.

        This method moves all SV objects from the database to the Bullet
        engine. Then it defers to Bullet for the physics update. Finally, it
        replaces the SV fields with the user specified values (only applies if
        the user called 'setBodyState') and writes the results back to the
        database.

        :param float dt: time step in seconds.
        :param int maxsteps: maximum number of sub-steps to simulate for one
                             ``dt`` update.
        """
        # Read queued commands and update the local object cache accordingly.
        with util.Timeit('Leonard:1.1  processCmdQueue') as timeit:
            self.processCommandQueue()

        # Update the constraint cache in our local Igor instance.
        self.igor.updateLocalCache()
        allConstraints = self.igor.getAllConstraints().data

        # Compute the collision sets.
        with util.Timeit('Leonard:1.2  CCS') as timeit:
            ret = self.igor.uniquePairs()
            if not ret.ok:
                return
            uniquePairs = ret.data

            ret = getFinalCollisionSets(
                uniquePairs, self.allObjects, self.allAABBs)
            if not ret.ok:
                return
            collSets = ret.data
            del ret, uniquePairs

        # Log the number of created collision sets.
        util.logMetricQty('#CollSets', len(collSets))

        # Put each collision set into its own Work Package.
        with util.Timeit('Leonard:1.3  CreateWPs') as timeit:
            all_WPs = {}
            for subset in collSets:
                # Compile the Work Package. Skip this physics step altogether
                # if an error occurs.
                ret = self.createWorkPackage(list(subset), dt, maxsteps)
                if not ret.ok:
                    self.logit.error(ret.msg)
                    return
                all_WPs[ret.data['wpid']] = ret.data

        with util.Timeit('Leonard:1.4  WPSendRecv') as timeit:
            wpIdx = 0
            worklist = list(all_WPs.keys())
            while True:
                # Wait until Worker contacts us. The message usually contains a
                # processed Work Package. However, it may also be empty, most
                # likely because the Worker has not received a Work Package
                # from us yet.
                msg = self.sock.recv()
                if msg != b'':
                    # Unpickle the message.
                    msg = pickle.loads(msg)
                    wpid = msg['wpid']

                    # Ignore the message if its Work Package is not pending
                    # anymore (most likely because multiple Workers processed
                    # the same Work Package and the other Worker has already
                    # returned it).
                    if wpid in all_WPs:
                        self.updateLocalCache(msg['wpdata'])

                        # Decrement the Work Package Index if the wpIdx counter
                        # is already past that work package. This simply
                        # ensures that no WP is skipped simply because the
                        # queue has shrunk.
                        if worklist.index(wpid) < wpIdx:
                            wpIdx -= 1

                        # Remove the WP from the work list and the WP cache.
                        worklist.remove(wpid)
                        del all_WPs[wpid]

                # Send an empty message to the Worker if no Work Packages are
                # pending anymore. This empty message is important to avoid
                # breaking the REQ/REP pattern enforced by ZeroMQ.
                if len(all_WPs) == 0:
                    self.sock.send(b'')
                    break

                # Pick the next pending Work Package and increment the index.
                if wpIdx >= len(worklist):
                    wpIdx = 0
                wp = all_WPs[worklist[wpIdx]]
                wpIdx += 1

                # Send the Work Package to the Worker.
                self.sock.send(pickle.dumps(wp))

        # Synchronise the local cache back to the database.
        with util.Timeit('Leonard:1.5  syncObjects') as timeit:
            self.syncObjects(writeconcern=False)

    @typecheck
    def createWorkPackage(self, objIDs: (tuple, list),
                          dt: (int, float), maxsteps: int):
        """
        Create a new Work Package (WP) and return its ID.

        The Work Package will not be returned but uploaded to the DB directly.

        A Work Package carries the necessary information for another rigid body
        physics steps. The Worker can thus start its work immediately, at least
        in terms of the rigid bodies (it may still want to incorporate other
        data like grid forces, but that is up to the Worker implementation).

        The ``dt`` and ``maxsteps`` arguments are for the underlying physics
        engine.

        .. note::
           A work package contains only the objIDs but not their SV. The
           ``getNextWorkPackage`` function takes care of compiling this
           information.

        :param iterable objIDs: list of object IDs in the new work package.
        :param float dt: time step for this work package.
        :param int maxsteps: number of sub-steps for the time step.
        :return: Work package ID
        :rtype: int
        """
        # Sanity check.
        if len(objIDs) == 0:
            return RetVal(False, 'Work package is empty', None)

        # Compile the Body States and forces into a list of ``WPData`` tuples.
        try:
            wpdata = []
            for objID in objIDs:
                sv = self.allObjects[objID]
                force, torque = self.totalForceAndTorque(objID)
                wpdata.append(WPData(objID, sv, force, torque))
        except KeyError as err:
            return RetVal(False, 'Cannot compile WP', None)

        # Query all constraints.
        constraints = self.igor.getConstraints(objIDs).data

        # Form the content of the Work Package as it will appear in the DB.
        data = {'wpid': self.wpid_counter,
                'wpmeta': (self.wpid_counter, dt, maxsteps),
                'wpdata': wpdata,
                'wpconstraints': constraints,
                'ts': None}
        self.wpid_counter += 1
        return RetVal(True, None, data)

    def updateLocalCache(self, wpdata):
        """
        Copy every object from ``wpdata`` to the local cache.

        The ``wpdata`` argument is a list of (objID, sv) tuples.

        The implicit assumption of this method is that ``wpdata`` is the
        output of ``computePhysicsForWorkPackage`` from a Worker.

        :param list wpdata: Content of Work Packge as returned by Workers.
        """
        # Reset force and torque for all objects in the WP, and overwrite
        # the old Body States with the new one from the processed WP.
        for (objID, sv) in wpdata:
            self.allObjects[objID] = _RigidBodyState(*sv)


class LeonardWorkerZeroMQ(config.AzraelProcess):
    """
    Dedicated Worker to process Work Packages.

    :param int workerID: the ID of this worker.
    :param int stepsUntilQuit: Worker will restart after this many steps.
    """
    def __init__(self, workerID, stepsUntilQuit: int):
        super().__init__()
        self.workerID = workerID

        # After ``stepsUntilQuit`` this Worker will spawn a new Worker with the
        # same ID and quit.
        assert stepsUntilQuit > 0
        self.stepsUntilQuit = stepsUntilQuit

        # Instantiate a Bullet engine.
        engine = azrael.bullet_api.PyBulletDynamicsWorld
        self.bullet = engine(self.workerID)

    def getGridForces(self, idPos: dict):
        """
        Return dictionary of force values for every object in ``idPos``.

        The ``idPos`` argument is a {objID_1: sv_1, objID_2, sv_2, ...}
        dictionary.

        The returned dictionary has the same keys as ``idPos``.

        :param dict idPos: dictionary with objIDs and corresponding SVs.
        :return dict: {objID_k: force_k}
        """
        # Convenience.
        vg = azrael.vectorgrid

        # Extract the keys and values in the same order.
        objIDs = list(idPos.keys())
        positions = [idPos[_] for _ in objIDs]

        # Query the grid values at all positions.
        ret = vg.getValues('force', positions)
        if not ret.ok:
            return RetVal(False, ret.msg, None)

        # Overwrite the default values with whatever ...
        gridForces = {objID: val for objID, val in zip(objIDs, ret.data)}
        return RetVal(True, None, gridForces)

    def computePhysicsForWorkPackage(self, wp):
        """
        Compute a physics steps for all objects in ``wp``.

        The output of this method is matched to the ``updateLocalCache`` in
        Leonard itself.

        :param dict wp: Work Package content from ``createWorkPackage``.
        :return dict: {'wpdata': list_of_SVs, 'wpid': wpid}
        """
        worklist, meta = wp['wpdata'], WPMeta(*wp['wpmeta'])
        constraints = wp['wpconstraints']

        # Log the number of collision-sets in the current Work Package.
        util.logMetricQty('Engine_{}'.format(self.workerID), len(worklist))

        # Convenience: the WPData elements make the code more readable.
        worklist = [WPData(*_) for _ in worklist]

        # Convenience.
        applyForceAndTorque = self.bullet.applyForceAndTorque
        setRigidBodyData = self.bullet.setRigidBodyData

        # Add every object to the Bullet engine and set the force/torque.
        with util.Timeit('Worker:1.1.0  applyforce') as timeit:
            with util.Timeit('Worker:1.1.1   grid') as timeit:
                # Fetch the forces for all object positions.
                idPos = {_.aid: _.sv.position for _ in worklist}
                ret = self.getGridForces(idPos)
                if not ret.ok:
                    self.logit.info(ret.msg)
                    z = np.float64(0)
                    gridForces = {_: z for _ in idPos}
                else:
                    gridForces = ret.data
                del ret, idPos

            with util.Timeit('Worker:1.1.1   updateGeo') as timeit:
                for obj in worklist:
                    # Update the object in Bullet and apply the force/torque.
                    setRigidBodyData(obj.aid, obj.sv)

            with util.Timeit('Worker:1.1.1   updateForce') as timeit:
                for obj in worklist:
                    # Add the force defined on the 'force' grid.
                    force = obj.force + gridForces[obj.aid]
                    applyForceAndTorque(obj.aid, force, obj.torque)

        # Apply all constraints. Log any errors but ignore them otherwise as
        # they are harmless (simply means no constraints were applied).
        ret = self.bullet.setConstraints(constraints)
        if not ret.ok:
            self.logit.warning(ret.msg)

        # Tell Bullet to advance the simulation for all objects in the
        # current work list.
        with util.Timeit('Worker:1.2.0  compute') as timeit:
            IDs = [_.aid for _ in worklist]
            self.bullet.compute(IDs, meta.dt, meta.maxsteps)

            # Remove all constraints.
            self.bullet.clearAllConstraints()

        with util.Timeit('Worker:1.3.0  fetchFromBullet') as timeit:
            # Retrieve the objects from Bullet again and update them in the DB.
            out = []
            for obj in worklist:
                ret = self.bullet.getRigidBodyData(obj.aid)
                sv = ret.data
                if not ret.ok:
                    # Something went wrong. Reuse the old SV.
                    sv = obj.sv
                    self.logit.error('Unable to get all objects from Bullet')
                out.append((obj.aid, sv))

        # Return the updated WP data.
        return {'wpid': meta.wpid, 'wpdata': out}

    @typecheck
    def run(self):
        """
        Wait for Work Packages, process them, and return the results.
        """
        # Call `run` method of `AzraelProcess` base class.
        super().run()

        try:
            # Setup ZeroMQ.
            ctx = zmq.Context()
            sock = ctx.socket(zmq.REQ)
            sock.connect(config.addr_leonard_repreq)
            self.logit.info('Worker {} connected'.format(self.workerID))

            # Contact Leonard with an empty payload.
            sock.send(b'')

            # Wait for messages from Leonard. If they contain a WP then process
            # it and return the result, otherwise reply with an empty message.
            numSteps = 0
            suq = self.stepsUntilQuit
            while numSteps < suq:
                # Wait for the next message.
                msg = sock.recv()

                # If Leonard did not send a Work Package (probably because it
                # does not have one right now) then wait for a short time
                # before asking again to avoid spamming the network.
                if msg == b'':
                    time.sleep(0.003)
                    sock.send(b'')
                    continue

                # Unpickle the Work Package.
                wpdata = pickle.loads(msg)

                # Process the Work Package.
                with util.Timeit('Worker:1.0.0 WPTotal') as timeit:
                    wpdata = self.computePhysicsForWorkPackage(wpdata)

                # Pack up the Work Package and send it back to Leonard.
                sock.send(pickle.dumps(wpdata))

                # Count the number of Work Packages we have processed.
                numSteps += 1

            # Log a last status message before terminating.
            self.logit.info('Worker {} terminated itself after {} steps'
                            .format(self.workerID, numSteps))
        except KeyboardInterrupt:
            print('Aborted Worker {}'.format(self.workerID))

        # Terminate.
        sock.close(linger=0)
        ctx.destroy()
        print('Worker {} exited cleanly'.format(self.workerID))


class WorkerManager(config.AzraelProcess):
    """
    Launch Worker processes and restart them as necessary.

    This class merely launches the inital set of workers and periodically
    checks if any have died. If so, it joins these processes and replaces it
    with a new Worker that has the same ID.

    :param int numWorker: number of Workers processes to spawn.
    :param int minSteps: see Worker
    :param int maxSteps: see Worker
    :param class workerCls: the class to instantiate.
    """
    def __init__(self, numWorkers: int, minSteps: int, maxSteps: int,
                 workerCls):
        super().__init__()

        # Sanity checks.
        assert numWorkers > 0
        assert isinstance(minSteps, int)
        assert isinstance(maxSteps, int)
        assert 0 < minSteps <= maxSteps

        # Backup the arguments.
        self.numWorkers = numWorkers
        self.workerCls = workerCls
        self.minSteps, self.maxSteps = minSteps, maxSteps

        # Initialise the list of workers.
        self.workers = []

    def _run(self):
        """
        Start the initial collection of Workers and ensure they remain alive.
        """
        # Spawn the initial collection of Workers.
        delta = self.maxSteps - self.minSteps
        for ii in range(self.numWorkers):
            # Random number in [minSteps, maxSteps]. The process will
            # automatically terminate after `suq` steps.
            suq = self.minSteps + int(np.random.rand() * delta)

            # Instantiate the process and add it to the list.
            self.workers.append(self.workerCls(ii + 1, suq))
            self.workers[-1].start()

        # Periodically monitor the processes and restart any that have died.
        while True:
            # Only check once a second.
            time.sleep(1)
            for workerID, proc in enumerate(self.workers):
                # Skip current process if it is still running.
                if proc.is_alive():
                    continue

                # Process has died --> join it to clear up the process table.
                proc.join()

                # Create a new Worker with the same ID but a (possibly)
                # different number of steps after which it must terminate.
                suq = self.minSteps + int(np.random.rand() * delta)
                proc = self.workerCls(workerID, suq)
                proc.start()
                self.workers[workerID] = proc
                self.logit.info('Restarted Worker {}'.format(workerID))

    def run(self):
        """
        Wrapper around ``_run`` to intercept SIGTERM.
        """
        # Call `run` method of `AzraelProcess` base class.
        super().run()

        try:
            self._run()
        except KeyboardInterrupt:
            print('Worker Manager was aborted')

        # Terminate all workers.
        for w in self.workers:
            if w is not None and w.is_alive():
                w.terminate()
                w.join()
        print('Worker manager finished')

{'fraction_non_alphanumeric': 0.08296806512210395, 'fraction_numerical': 0.004620732839591475, 'mean_word_length': 1.9481205639899424, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 25, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '183991', 'n_tokens_mistral': 26172, 'n_tokens_neox': 24838, 'n_words': 9267}