class ForceField(simtk.openmm.app.ForceField):

    def registerBodyTemplate(self, name, residue, pattern=None):
    """
        Register a new rigid body template.

        Parameters
        ----------
        name : string
            A name for the rigid body template being registered.
        residue : string
            The name of the residue template to which this body belongs.
        pattern : string, optional, default=None
            A string containing a regular expression for selecting the atoms in residue which will
            form the body. A list of atom names can be passed as '(name1|name2|...)', for instance.
            Only actual atoms will be considered (i.e. virtual sites will be ignored). If patttern
            is None, then the whole residue will form the rigid body.
      """

      def createSystem(self, topology, **kwargs):
      """
          Construct an OpenMM System representing a Topology with this force field. This extends
          the original method by creating rigid bodies from previously registered body templates.

          Parameters
          ----------
          topology : Topology
              The Topology for which to create a System.
          kwargs : multiple types
              All keyword arguments of the original `ForceField.createSystem` method, plus:
          mergeList : list(list(int)), optional, default=None
              A nested list of indices of rigid bodies to be merged together after the templates
              are applied. If this is None, then no merging will occur.
          removeConstraints : bool, optional, default=False
              If true, every constraint involving at least one atom that belongs to a rigid body
              will be silently removed from the system.
          removeForces : bool, optional, default=False
              If true, all interactions whose all involved atoms belong to the same rigid body
              will be silently removed from the system.

          Returns
          -------
          system : System
              The newly created System.
          bodyIndices : list(int)
              A list containing the index of the rigid body to which every atom belongs (note:
              index=0 for free atoms).
      """

class RigidBodyIntegrator(simtk.openmm.Integrator):

    def __init__(self, stepSize, bodyIndices):
    """
        Create a Rigid Body Integrator.

        Parameters
        ----------
        stepSize : double
            The step size with which to integrate the system (in picoseconds).
        bodyIndices : list(int)
            A list containing the index of the rigid body to which every atom belongs (note:
            index=0 for free atoms).
    """

    def setRotationMode(mode):
    """
        Defines the mode of rotational move integration.

        Parameters
        ----------
        mode : int
            The mode of rotation, where mode=0 consisting in using an exact solution for rotations
            spanning the whole time step, while mode=n consists in using a NO-SQUISH solution with
            the time step split into n small steps.
    """

    def step(self, steps):
    """
       	Advance a simulation through time by taking a series of time steps.

        Parameters
        ----------
        steps : int
            The number of time steps to be taken.
    """

class StateDataReporter(simtk.openmm.app.StateDataReporter):

    def __init__(self, *args, **kwargs):
    """
        Create a State Data Reporter.

        Parameters
        ----------
        args : multiple types
            All positional arguments of the original `StateDataReporter` constructor.
        kwargs : multiple types
            All keyword arguments of the original `StateDataReporter` constructor, plus:
        translationalEnergy : bool, optional, default=False
            Whether to write the translational part of the kinetic energy to the file.
        rotationalEnergy : bool, optional, default=False
            Whether to write the rotational part of the kinetic energy to the file.
    """

import simtk.openmm as mm
from simtk.openmm import app
import rigidbodyplugin as rbmm
pdb = app.PDBFile('config.pdb')
forcefield = rbmm.ForceField('model.xml')
forcefield.registerBodyTemplate('water', 'HOH')
(system, bodies) = forcefield.createSystem(pdb.topology, nonbondedMethod=app.PME,
                                           removeConstraints=True, removeForces=True)
integrator = rbmm.RigidBodyIntegrator(1.0*unit.femtoseconds, bodies)
integrator.setRotationMode(0)
reporter = rbmm.StateDataReporter(stdout, 1, step=True, temperature=True,
                                  translationalEnergy=True, rotationalEnergy=True)
simulation = app.Simulation(pdb.topology, system, integrator,
                            platform=mm.Platform.getPlatformByName('CUDA'),
                            properties={'Precision': 'mixed'})
simulation.context.setPositions(pdb.positions)
simulation.context.setVelocitiesToTemperature(300*unit.kelvin)
simulation.reporters.append(reporter)
simulation.step(1000)