#!/usr/local/bin/env python

# ==============================================================================
# FILE DOCSTRING
# ==============================================================================

"""

Restraints
==========

Automated selection and imposition of receptor-ligand restraints for absolute
alchemical binding free energy calculations, along with computation of the
standard state correction.


"""

# ==============================================================================
# GLOBAL IMPORTS
# ==============================================================================

import abc
import random
import inspect
import logging
import functools
import itertools

import numpy as np
import scipy.integrate
import mdtraj as md
import openmmtools as mmtools
from simtk import openmm, unit

from . import pipeline
from .utils import methoddispatch, generate_development_feature

from typing import Set, Iterable

logger = logging.getLogger(__name__)
ABC = abc.ABC


# ==============================================================================
# MODULE CONSTANTS
# ==============================================================================

V0 = 1660.53928 * unit.angstroms**3  # standard state volume


# ==============================================================================
# CUSTOM EXCEPTIONS
# ==============================================================================

class RestraintStateError(mmtools.states.ComposableStateError):
    """Error raised by an :class:`RestraintState`."""
    pass


class RestraintParameterError(Exception):
    """Error raised by a :class:`ReceptorLigandRestraint`."""
    pass


# ==============================================================================
# DEVELOPMENT FEATURES
# ==============================================================================

# Cannot simply test .__version__ since that does not get updated until release
OpenMM73 = generate_development_feature({"OpenMM >=7.3": hasattr(openmm, "RMSDForce")})


# ==============================================================================
# Dispatch appropriate restraint type from registered restraint classes
# ==============================================================================

def available_restraint_classes():
    """
    Return all available restraint classes.

    Returns
    -------
    restraint_classes : dict of {str : class}
        ``restraint_classes[name]`` is the class corresponding to ``name``

    """
    restraint_subclasses = mmtools.utils.find_all_subclasses(ReceptorLigandRestraint,
                                                             discard_abstract=True)

    # Build an index of all names, ensuring there are no name collisions.
    available_restraints = dict()
    for cls in restraint_subclasses:
        cls_name = cls.__name__
        if cls_name in available_restraints:
            raise ValueError("More than one restraint subclass has the name '{}'.".format(cls_name))
        else:
            available_restraints[cls_name] = cls

    return available_restraints


def available_restraint_types():
    """
    List all available restraint types.

    Returns
    -------
    available_restraint_types : list of str
        List of names of available restraint classes

    """
    available_restraints = available_restraint_classes()
    return list(available_restraints.keys())


def create_restraint(restraint_type, **kwargs):
    """Factory of receptor-ligand restraint objects.

    Parameters
    ----------
    restraint_type : str
        Restraint type name matching a register (imported) subclass of :class:`ReceptorLigandRestraint`.
    kwargs
        Parameters to pass to the restraint constructor.

    """
    cls = mmtools.utils.find_subclass(ReceptorLigandRestraint, restraint_type)
    return cls(**kwargs)


# ==============================================================================
# ComposableState class to control the strength of restraints.
# ==============================================================================

class RestraintState(object):
    """
    The state of a restraint.

    A ``ComposableState`` controlling the strength of a restraint
    through its ``lambda_restraints`` property.

    Parameters
    ----------
    lambda_restraints : float
        The strength of the restraint. Must be between 0 and 1.

    Attributes
    ----------
    lambda_restraints

    Examples
    --------
    Create a system in a thermodynamic state.

    >>> from openmmtools import testsystems, states
    >>> system_container = testsystems.LysozymeImplicit()
    >>> system, positions = system_container.system, system_container.positions
    >>> thermodynamic_state = states.ThermodynamicState(system, 300*unit.kelvin)
    >>> sampler_state = states.SamplerState(positions)

    Identify ligand atoms. Topography automatically identify receptor atoms too.

    >>> from yank.yank import Topography
    >>> topography = Topography(system_container.topology, ligand_atoms=range(2603, 2621))

    Apply a Harmonic restraint between receptor and protein. Let the restraint
    automatically determine all the parameters.

    >>> restraint = Harmonic()
    >>> restraint.determine_missing_parameters(thermodynamic_state, sampler_state, topography)
    >>> restraint.restrain_state(thermodynamic_state)

    Create a ``RestraintState`` object to control the strength of the restraint.

    >>> restraint_state = RestraintState(lambda_restraints=1.0)

    ``RestraintState`` implements the ``IComposableState`` interface, so it can be
    used with ``CompoundThermodynamicState``.

    >>> compound_state = states.CompoundThermodynamicState(thermodynamic_state=thermodynamic_state,
    ...                                                    composable_states=[restraint_state])
    >>> compound_state.lambda_restraints
    1.0
    >>> integrator = openmm.VerletIntegrator(1.0*unit.femtosecond)
    >>> context = compound_state.create_context(integrator)
    >>> context.getParameter('lambda_restraints')
    1.0

    You can control the parameters in the OpenMM Context by setting the state's
    attributes. To To deactivate the restraint, set `lambda_restraints` to 0.0.

    >>> compound_state.lambda_restraints = 0.0
    >>> compound_state.apply_to_context(context)
    >>> context.getParameter('lambda_restraints')
    0.0

    """
    def __init__(self, lambda_restraints):
        self.lambda_restraints = lambda_restraints

    @property
    def lambda_restraints(self):
        """Float: the strength of the applied restraint (between 0 and 1 inclusive)."""
        return self._lambda_restraints

    @lambda_restraints.setter
    def lambda_restraints(self, value):
        assert 0.0 <= value <= 1.0
        self._lambda_restraints = float(value)

    def apply_to_system(self, system):
        """
        Set the strength of the system's restraint to this.

        System is updated in-place

        Parameters
        ----------
        system : simtk.openmm.System
            The system to modify.

        Raises
        ------
        RestraintStateError
            If the system does not have any ``CustomForce`` with a
            ``lambda_restraint`` global parameter.

        """
        # Set lambda_restraints in all forces that have it.
        for force, parameter_id in self._get_system_forces_parameters(system):
            force.setGlobalParameterDefaultValue(parameter_id, self._lambda_restraints)

    def check_system_consistency(self, system):
        """
        Check if the system's restraint is in this restraint state.

        It raises a :class:`RestraintStateError` if the restraint is not consistent
        with the state.

        Parameters
        ----------
        system : simtk.openmm.System
            The system with the restraint to test.

        Raises
        ------
        RestraintStateError
            If the system is not consistent with this state.

        """
        # Set lambda_restraints in all forces that have it.
        for force, parameter_id in self._get_system_forces_parameters(system):
            force_lambda = force.getGlobalParameterDefaultValue(parameter_id)
            if force_lambda != self.lambda_restraints:
                err_msg = 'Consistency check failed: system {}, state {}'
                raise RestraintStateError(err_msg.format(force_lambda, self._lambda_restraints))

    def apply_to_context(self, context):
        """Put the restraint in the `Context` into this state.

        Parameters
        ----------
        context : simtk.openmm.Context
            The context to set.

        Raises
        ------
        RestraintStateError
            If the context does not have the required lambda global variables.

        """
        try:
            context.setParameter('lambda_restraints', self._lambda_restraints)
        except Exception:
            raise RestraintStateError('The context does not have a restraint.')

    def _standardize_system(self, system):
        """Standardize the given system.

        Set lambda_restraints of the system to 1.0.

        Parameters
        ----------
        system : simtk.openmm.System
            The system to standardize.

        Raises
        ------
        RestraintStateError
            If the system is not consistent with this state.

        """
        # Set lambda_restraints to 1.0 in all forces that have it.
        for force, parameter_id in self._get_system_forces_parameters(system):
            force.setGlobalParameterDefaultValue(parameter_id, 1.0)

    def _on_setattr(self, standard_system, attribute_name):
        """Check if the standard system needs changes after a state attribute is set.

        Parameters
        ----------
        standard_system : simtk.openmm.System
            The standard system before setting the attribute.
        attribute_name : str
            The name of the attribute that has just been set or retrieved.

        Returns
        -------
        need_changes : bool
            True if the standard system has to be updated, False if no change
            occurred.

        """
        # There are no attributes that can be set that can alter the standard system.
        return False

    def _find_force_groups_to_update(self, context, current_context_state, memo):
        """Find the force groups whose energy must be recomputed after applying self.

        Parameters
        ----------
        context : Context
            The context, currently in `current_context_state`, that will
            be moved to this state.
        current_context_state : ThermodynamicState
            The full thermodynamic state of the given context. This is
            guaranteed to be compatible with self.
        memo : dict
            A dictionary that can be used by the state for memoization
            to speed up consecutive calls on the same context.

        Returns
        -------
        force_groups_to_update : set of int
            The indices of the force groups whose energy must be computed
            again after applying this state, assuming the context to be in
            `current_context_state`.
        """
        # Check if lambda_restraints will change.
        if self.lambda_restraints == current_context_state.lambda_restraints:
            return set()

        # Update memo if this is the first call for this context.
        if len(memo) == 0:
            system = context.getSystem()
            for force, _ in self._get_system_forces_parameters(system):
                memo['lambda_restraints'] = force.getForceGroup()
        return {memo['lambda_restraints']}

    @staticmethod
    def _get_system_forces_parameters(system):
        """Yields the system's forces having a ``lambda_restraints`` parameter.

        Yields
        ------
        A tuple force, ``parameter_index`` for each force with ``lambda_restraints``.

        """
        found_restraint = False

        # Retrieve all the forces with global supported parameters.
        for force_index in range(system.getNumForces()):
            force = system.getForce(force_index)
            try:
                n_global_parameters = force.getNumGlobalParameters()
            except AttributeError:
                continue
            for parameter_id in range(n_global_parameters):
                parameter_name = force.getGlobalParameterName(parameter_id)
                if parameter_name == 'lambda_restraints':
                    found_restraint = True
                    yield force, parameter_id
                    break

        # Raise error if the system doesn't have a restraint.
        if found_restraint is False:
            raise RestraintStateError('The system does not have a restraint.')

    def __getstate__(self):
        return dict(lambda_restraints=self._lambda_restraints)

    def __setstate__(self, serialization):
        self.lambda_restraints = serialization['lambda_restraints']


# ==============================================================================
# Atom selector tool
# ==============================================================================

class _AtomSelector(object):
    """
    Helper class to select atoms from semi-arbitrary selections based on a topography
    """

    def __init__(self, topography):
        self.topography = topography

    def compute_atom_intersect(self, input_atoms, topography_key: str, *additional_sets: Iterable[int]) -> Set[int]:

        """
        Compute the intersect of atoms formed from a given input, topography key to reference, and
        any additional sets for cross reference

        Parameters
        ----------
        input_atoms : str, iterable of int, or None
            Atom-like selection which can accept a :func:`yank.Topography.select`, or a sequence of ints, or None
            When given None, only the Topography key and additional sets are used
        topography_key : str
            Key in the :class:`yank.Topography` which is used to initially sub-select atoms
        additional_sets : set of int or iterable of int
            Any additional sets to cross reference

        Returns
        -------
        atom_intersect : set
            Set of atoms intersecting the input, the topography key, and the additional sets
        """

        topography = self.topography
        topography_set = set(getattr(topography, topography_key))
        # Ensure additions are sets
        additional_sets = [set(additional_set) for additional_set in additional_sets]
        if len(additional_sets) == 0:
            # Case no sets were provided
            additional_intersect = topography_set
        else:
            additional_intersect = set.intersection(*additional_sets)

        @functools.singledispatch
        def compute_atom_set(passed_atoms):
            """Helper function for doing set operations on heavy ligand atoms of all other types"""
            input_set = set(passed_atoms)
            intersect_set = input_set & additional_intersect & topography_set
            if intersect_set != input_set:
                return intersect_set
            else:
                # This ensures if no changes are made to the set, then passed atoms are returned unmodied
                return passed_atoms

        @compute_atom_set.register(type(None))
        def compute_atom_none(_):
            """Helper for None type parsing"""
            return topography_set & additional_intersect

        @compute_atom_set.register(str)
        def compute_atom_str(input_string):
            """Helper for string parsing"""
            output = topography.select(input_string, as_set=False)  # Preserve order
            set_output = set(output)
            # Ensure the selection is in the correct set
            set_combined = set_output & topography_set & additional_intersect
            final_output = [particle for particle in output if particle in set_combined]
            # Force output to be a normal int, don't need to worry about floats at this point, there should not be any
            # If they come out as np.int64's, OpenMM complains
            return [*map(int, final_output)]

        return compute_atom_set(input_atoms)


# ==============================================================================
# Base class for receptor-ligand restraints.
# ==============================================================================

class ReceptorLigandRestraint(ABC):
    """
    A restraint preventing a ligand from drifting too far from its receptor.

    With replica exchange simulations, keeping the ligand close to the binding
    pocket can enhance mixing between the interacting and the decoupled state.
    This should be always used in implicit simulation, where there are no periodic
    boundary conditions.

    This restraint strength is controlled by a global context parameter called
    ``lambda_restraints``. You can easily control this variable through the
    ``RestraintState`` object.

    Notes
    -----
    Creating a subclass requires the following:

        1. Implement a constructor. Optionally this can leave all or a subset of
        the restraint parameters undefined. In this case, you need to provide
        an implementation of :func:`determine_missing_parameters`.

        2. Implement :func:`restrain_state` that add the restrain ``Force`` to the state's
        `System`.

        3. Implement :func:`get_standard_state_correction` to return standard state correction.

        4. Optionally, implement :func:`determine_missing_parameters` to fill in
        the parameters left undefined in the constructor.

    """

    @abc.abstractmethod
    def restrain_state(self, thermodynamic_state):
        """Add the restraint force to the state's `System`.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state holding the system to modify.

        """
        pass

    @abc.abstractmethod
    def get_standard_state_correction(self, thermodynamic_state):
        """Return the standard state correction.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state.

        """
        pass

    def determine_missing_parameters(self, thermodynamic_state, sampler_state, topography):
        """
        Automatically choose undefined parameters.

        Optionally, a :class:`ReceptorLigandRestraint` can support the automatic
        determination of all or a subset of the parameters that can be
        left undefined in the constructor, making implementation of this method optional.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynmaic state to inspect
        sampler_state : openmmtools.states.SamplerState
            The sampler state holding the positions of all atoms.
        topography : yank.Topography
            The topography with labeled receptor and ligand atoms.

        """
        raise NotImplementedError('{} does not support automatic determination of the '
                                  'restraint parameters'.format(self.__class__.__name__))

    @classmethod
    def _add_force_in_separate_group(cls, system, restraint_force):
        """Add the force to the System in a separate force group when possible."""
        # OpenMM supports a maximum of 32 force groups.
        available_force_groups = set(range(32))
        for force in system.getForces():
            available_force_groups.discard(force.getForceGroup())

        # If the System is full, just separate the force from nonbonded interactions.
        if len(available_force_groups) == 0:
            _, nonbonded_force = mmtools.forces.find_forces(system, openmm.NonbondedForce,
                                                            only_one=True)
            available_force_groups = set(range(32))
            available_force_groups.discard(nonbonded_force.getForceGroup())

        restraint_force.setForceGroup(min(available_force_groups))
        system.addForce(restraint_force)

    @property
    def _excluded_init_parameters(self):
        """
        List of excluded parameters from the :func:`__init__` call to ensure all non-atom selection
        parameters are defined
        """
        return ['self', 'restrained_receptor_atoms', 'restrained_ligand_atoms']

    @property
    def _parameters(self):
        """dict: restraint parameters in dict forms."""
        argspec = inspect.getfullargspec(self.__init__)
        parameter_names = argspec.args

        # Exclude non-parameters arguments.
        for exclusion in self._excluded_init_parameters:
            parameter_names.remove(exclusion)

        # Retrieve and store options.
        parameters = {parameter_name: getattr(self, parameter_name)
                      for parameter_name in parameter_names}
        return parameters


class _RestrainedAtomsProperty(object):
    """
    Descriptor of restrained atoms.

    Casts generic iterables of ints into lists.
    """

    def __init__(self, atoms_type):
        self._atoms_type = atoms_type

    @property
    def _attribute_name(self):
        """Name of the internally stored variable (read-only)."""
        return '_restrained_' + self._atoms_type + '_atoms'

    def __get__(self, instance, owner_class=None):
        return getattr(instance, self._attribute_name)

    def __set__(self, instance, new_restrained_atoms):
        # If we set the restrained attributes to None, no reason to check things.
        if new_restrained_atoms is not None:
            new_restrained_atoms = self._validate_atoms(new_restrained_atoms)
        setattr(instance, self._attribute_name, new_restrained_atoms)

    @methoddispatch
    def _validate_atoms(self, restrained_atoms):
        """Casts a generic iterable of ints into a list to support concatenation."""
        try:
            restrained_atoms = restrained_atoms.tolist()
        except AttributeError:
            restrained_atoms = list(restrained_atoms)
        return restrained_atoms


# ==============================================================================
# Base class for radially-symmetric receptor-ligand restraints.
# ==============================================================================

class RadiallySymmetricRestraint(ReceptorLigandRestraint):
    """
    Base class for radially-symmetric restraints between ligand and protein.

    The restraint is applied between the centroids of two groups of atoms
    that belong to the receptor and the ligand respectively. The centroids
    are determined by a mass-weighted average of the group particles positions.
    The restraint strength is controlled by a global context parameter called
    'lambda_restraints'.

    With OpenCL, groups with more than 1 atom are supported only on 64bit
    platforms.

    The class allows the restrained atoms to be temporarily undefined, but in
    this case, :func:`determine_missing_parameters` must be called before using
    the restraint.

    Parameters
    ----------
    restrained_receptor_atoms : iterable of int, int, or str, optional
        The indices of the receptor atoms to restrain, an MDTraj DSL expression, any other
        :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Selection <yank.Topography.select>`.
        This can temporarily be left undefined, but :func:`determine_missing_parameters`
        must be called before using the Restraint object. The same if a DSL
        expression or Topography selection is provided (default is None).
    restrained_ligand_atoms : iterable of int, int, or str, optional
        The indices of the ligand atoms to restrain, an MDTraj DSL expression, or a
        :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Selection <yank.Topography.select>`.
        This can temporarily be left undefined, but :func:`determine_missing_parameters`
        must be called before using the Restraint object. The same if a DSL
        expression or Topography selection is provided (default is None).

    Attributes
    ----------
    restrained_receptor_atoms : list of int, str, or None
        The indices of the receptor atoms to restrain, an MDTraj selection string, or a Topography selection
        string.
    restrained_ligand_atoms : list of int, str, None
        The indices of the receptor atoms to restrain, an MDTraj selection string, or a Topography selection
        string.

    Notes
    -----
    To create a subclass, follow these steps:

        1. Implement the :func:`_create_restraint_force` returning the
        force object modeling the restraint.

        2. Implement the property :func:`_are_restraint_parameters_defined`.

        3. Optionally, you can overwrite the :func:`_determine_restraint_parameters`
        method to automatically determine these parameters from the atoms positions.

    """
    def __init__(self, restrained_receptor_atoms=None, restrained_ligand_atoms=None):
        self.restrained_receptor_atoms = restrained_receptor_atoms
        self.restrained_ligand_atoms = restrained_ligand_atoms

    # -------------------------------------------------------------------------
    # Public properties.
    # -------------------------------------------------------------------------

    class _RadiallySymmetricRestrainedAtomsProperty(_RestrainedAtomsProperty):
        """
        Descriptor of restrained atoms.

        Extends `_RestrainedAtomsProperty` to handle single integers and strings.
        """
        _CENTROID_COMPUTE_STRING = ("You are specifying {} {} atoms, "
                                    "the final atoms will be chosen as the centroid of this set.")

        @methoddispatch
        def _validate_atoms(self, restrained_atoms):
            restrained_atoms = super()._validate_atoms(restrained_atoms)
            if len(restrained_atoms) > 1:
                logger.debug(self._CENTROID_COMPUTE_STRING.format("more than one", self._atoms_type))
            return restrained_atoms

        @_validate_atoms.register(str)
        def _validate_atoms_string(self, restrained_atoms):
            warn_string = self._CENTROID_COMPUTE_STRING.format("a string for", self._atoms_type)
            warn_string += 'but you MUST run "determine_missing_parameters" to process the string'
            logger.warning(warn_string)
            return restrained_atoms

        @_validate_atoms.register(int)
        def _validate_atoms_int(self, restrained_atoms):
            return [restrained_atoms]

    restrained_receptor_atoms = _RadiallySymmetricRestrainedAtomsProperty('receptor')
    restrained_ligand_atoms = _RadiallySymmetricRestrainedAtomsProperty('ligand')

    # -------------------------------------------------------------------------
    # Public methods.
    # -------------------------------------------------------------------------

    def restrain_state(self, thermodynamic_state):
        """Add the restraining Force(s) to the thermodynamic state's system.

        All the parameters must be defined at this point. An exception is
        raised if they are not.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state holding the system to modify.

        Raises
        ------
        RestraintParameterError
            If the restraint has undefined parameters.

        """
        # Check that restrained atoms are defined.
        if not self._are_restrained_atoms_defined:
            raise RestraintParameterError('Restraint {}: Undefined restrained '
                                          'atoms.'.format(self.__class__.__name__))

        # Create restraint force.
        restraint_force = self._get_restraint_force(self.restrained_receptor_atoms,
                                                    self.restrained_ligand_atoms)

        # Set periodic conditions on the force if necessary.
        restraint_force.setUsesPeriodicBoundaryConditions(thermodynamic_state.is_periodic)

        # Get a copy of the system of the ThermodynamicState, modify it and set it back.
        system = thermodynamic_state.system
        self._add_force_in_separate_group(system, restraint_force)
        thermodynamic_state.system = system

    def get_standard_state_correction(self, thermodynamic_state):
        """Return the standard state correction.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state.

        Returns
        -------
        correction : float
           The unit-less standard-state correction, in kT (at the
           temperature of the given thermodynamic state).

        """
        benchmark_id = 'Restraint {}: Computing standard state correction'.format(self.__class__.__name__)
        timer = mmtools.utils.Timer()
        timer.start(benchmark_id)

        restraint_force = self._get_restraint_force([0], [1])
        DeltaG = restraint_force.compute_standard_state_correction(thermodynamic_state, max_volume='system')
        logger.debug('Standard state correction: {:.3f} kT'.format(DeltaG))

        # Report elapsed time.
        timer.stop(benchmark_id)
        timer.report_timing()

        # Return standard state correction (in kT).
        return DeltaG

    def determine_missing_parameters(self, thermodynamic_state, sampler_state, topography):
        """Automatically determine missing parameters.

        If some parameters have been left undefined (i.e. the atoms to restrain
        or the restraint force parameters) this attempts to find them using the
        information in the states and the topography.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state.
        sampler_state : openmmtools.states.SamplerState, optional
            The sampler state holding the positions of all atoms.
        topography : yank.Topography, optional
            The topography with labeled receptor and ligand atoms.

        """

        # Determine restrained atoms, if needed.
        self._determine_restrained_atoms(sampler_state, topography)

        # Determine missing parameters. This is implemented in the subclass.
        self._determine_restraint_parameters(thermodynamic_state, sampler_state, topography)

    # -------------------------------------------------------------------------
    # Internal-usage: properties and methods for subclasses.
    # -------------------------------------------------------------------------

    @property
    @abc.abstractmethod
    def _are_restraint_parameters_defined(self):
        """bool: True if the restraint parameters are defined."""
        pass

    @abc.abstractmethod
    def _create_restraint_force(self, particles1, particles2):
        """Create a new restraint force between specified atoms.

        The property _are_restraint_parameters_defined is guaranteed to
        be True when this is called.

        Parameters
        ----------
        particles1 : list of int
            Indices of first group of atoms to restraint.
        particles2 : list of int
            Indices of second group of atoms to restraint.

        Returns
        -------
        force : simtk.openmm.Force
           The created restraint force.

        """
        pass

    def _determine_restraint_parameters(self, thermodynamic_state, sampler_state, topography):
        """Determine the missing bond parameters.

        Optionally, a subclass can implement this method to automatically
        define the bond parameters of the restraints from the information
        in the given states and topography. The default implementation just
        raises a NotImplemented error if `_are_restraint_parameters_defined`
        is False.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state.
        sampler_state : openmmtools.states.SamplerState
            The sampler state holding the positions of all atoms.
        topography : yank.Topography
            The topography with labeled receptor, ligand atoms, and any regions defined.

        """
        # Raise exception only if the subclass doesn't already defines parameters.
        if not self._are_restraint_parameters_defined:
            raise NotImplementedError('Restraint {} cannot automatically determine '
                                      'bond parameters.'.format(self.__class__.__name__))

    # -------------------------------------------------------------------------
    # Internal-usage
    # -------------------------------------------------------------------------

    @property
    def _are_restrained_atoms_defined(self):
        """Check if the restrained atoms are defined well enough to make a restraint"""
        for atoms in [self.restrained_receptor_atoms, self.restrained_ligand_atoms]:
            # Atoms should be a list or None at this point due to the _RestrainedAtomsProperty class
            if atoms is None or not (isinstance(atoms, list) and len(atoms) > 0):
                return False
        return True

    def _get_restraint_force(self, particles1, particles2):
        """Check that the parameters are defined before calling the force constructor."""
        # Check if restraint parameters have been defined.
        if self._are_restraint_parameters_defined is None:
            err_msg = 'Restraint {}: Undefined bond parameters.'.format(self.__class__.__name__)
            raise RestraintParameterError(err_msg)

        # Create restraint force.
        return self._create_restraint_force(particles1, particles2)

    def _determine_restrained_atoms(self, sampler_state, topography):
        """Determine the atoms to restrain.

        If the user has explicitly specified which atoms to restrained, this
        does nothing, otherwise it picks the centroid of the receptor and
        the centroid of the ligand as the two atoms to restrain.

        Parameters
        ----------
        sampler_state : openmmtools.states.SamplerState, optional
            The sampler state holding the positions of all atoms.
        topography : yank.Topography, optional
            The topography with labeled receptor, ligand atoms, and any regions defined.

        """
        debug_msg = ('Restraint {}: Automatically picked restrained '
                     '{{0}} atom: {{0}}'.format(self.__class__.__name__))

        # No need to determine parameters if atoms have been given.
        if self._are_restrained_atoms_defined:
            return

        # Shortcuts
        positions = sampler_state.positions

        # If receptor and ligand atoms are explicitly provided, use those.
        restrained_ligand_atoms = self.restrained_ligand_atoms
        restrained_receptor_atoms = self.restrained_receptor_atoms

        @functools.singledispatch
        def compute_atom_set(input_atoms, topography_key, mapping_function):
            """
            Helper function for doing set operations on generic atom types.
            mapping_function not used in the generic catch-all, but is used in the None register
            """
            # Ensure the input atoms are only part of the topography_key atoms. Make no changes if they are
            input_atoms_set = set(input_atoms)
            set_topography_atoms = set(getattr(topography, topography_key))
            intersect_set = input_atoms_set & set_topography_atoms
            if intersect_set != input_atoms_set:
                logger.warning("Some atoms specified by {0} were not actual {0}! "
                               "Atoms not part of {0} will be ignored.".format(topography_key))
                final_atoms = list(intersect_set)
            else:
                final_atoms = list(input_atoms)
            return final_atoms

        @compute_atom_set.register(type(None))
        def compute_atom_none(_, topography_key, mapping_function):
            """Helper for None type parsing"""
            # Can't use list() here since mapping function returns a single integer.
            atom_selection = [mapping_function(positions, getattr(topography, topography_key))]
            logger.debug(debug_msg.format(topography_key, atom_selection))
            return atom_selection

        @compute_atom_set.register(str)
        def compute_atom_str(input_string, topography_key, _):
            """Helper for string parsing"""
            selection = topography.select(input_string, as_set=True)
            selection_with_top = selection & set(getattr(topography, topography_key))
            # Force output to be a normal int, dont need to worry about floats at this point, there should not be any
            # If they come out as np.int64's, OpenMM complains
            return [*map(int, selection_with_top)]

        self.restrained_ligand_atoms = compute_atom_set(restrained_ligand_atoms,
                                                        'ligand_atoms',
                                                        self._closest_atom_to_centroid)
        self.restrained_receptor_atoms = compute_atom_set(restrained_receptor_atoms,
                                                          'receptor_atoms',
                                                          self._closest_atom_to_centroid)

    @staticmethod
    def _closest_atom_to_centroid(positions, indices=None, masses=None):
        """
        Identify the closest atom to the centroid of the given coordinate set.

        Parameters
        ----------
        positions : unit.Quantity of natoms x 3 with units compatible with nanometers
           positions of object to identify atom closes to centroid
        indices : list of int, optional, default=None
           List of atoms indices for which closest atom to centroid is to be computed.
        masses : simtk.unit.Quantity of natoms with units compatible with amu
           Masses of particles used to weight distance calculation, if not None (default: None)

        Returns
        -------
        closest_atom : int
           Index of atom closest to centroid of specified atoms.

        """
        if indices is not None:
            positions = positions[indices, :]

        # Get dimensionless positions.
        x_unit = positions.unit
        x = positions / x_unit

        # Determine number of atoms.
        natoms = x.shape[0]

        # Compute (natoms,1) array of normalized weights.
        w = np.ones([natoms, 1])
        if masses is not None:
            w = masses / masses.unit # (natoms,) array
            w = np.reshape(w, (natoms, 1))  # (natoms,1) array
        w /= w.sum()

        # Compute centroid (still in dimensionless units).
        centroid = (np.tile(w, (1, 3)) * x).sum(0)  # (3,) array

        # Compute distances from centroid.
        distances = np.sqrt(((x - np.tile(centroid, (natoms, 1)))**2).sum(1))  # distances[i] is the distance from the centroid to particle i

        # Determine closest atom.
        closest_atom = int(np.argmin(distances))

        if indices is not None:
            closest_atom = indices[closest_atom]

        return closest_atom


# ==============================================================================
# Harmonic protein-ligand restraint.
# ==============================================================================

class Harmonic(RadiallySymmetricRestraint):
    """Impose a single harmonic restraint between ligand and protein.

    This can be used to prevent the ligand from drifting too far from the
    protein in implicit solvent calculations or to keep the ligand close
    to the binding pocket in the decoupled states to increase mixing.

    The restraint is applied between the centroids of two groups of atoms
    that belong to the receptor and the ligand respectively. The centroids
    are determined by a mass-weighted average of the group particles positions.

    The energy expression of the restraint is given by

       ``E = lambda_restraints * (K/2)*r^2``

    where `K` is the spring constant, `r` is the distance between the
    two group centroids, and `lambda_restraints` is a scale factor that
    can be used to control the strength of the restraint. You can control
    ``lambda_restraints`` through :class:`RestraintState` class.

    The class supports automatic determination of the parameters left undefined or defined by strings
    in the constructor through :func:`determine_missing_parameters`.

    With OpenCL, groups with more than 1 atom are supported only on 64bit
    platforms.

    Parameters
    ----------
    spring_constant : simtk.unit.Quantity, optional
        The spring constant K (see energy expression above) in units compatible
        with joule/nanometer**2/mole (default is None).
    restrained_receptor_atoms : iterable of int, int, or str, optional
        The indices of the receptor atoms to restrain, an MDTraj DSL expression, or a
        :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Select String <yank.Topography.select>`.
        This can temporarily be left undefined, but ``determine_missing_parameters()``
        must be called before using the Restraint object. The same if a DSL
        expression or Topography region is provided (default is None).
    restrained_ligand_atoms : iterable of int, int, or str, optional
        The indices of the ligand atoms to restrain, an MDTraj DSL expression.
        or a :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Select String <yank.Topography.select>`.
        This can temporarily be left undefined, but ``determine_missing_parameters()``
        must be called before using the Restraint object. The same if a DSL
        expression or Topography region is provided (default is None).

    Attributes
    ----------
    restrained_receptor_atoms : list of int, str, or None
        The indices of the receptor atoms to restrain, an MDTraj selection string, or a Topography region selection
        string.
    restrained_ligand_atoms : list of int, str, or None
        The indices of the ligand atoms to restrain, an MDTraj selection string, or a Topography region selection
        string.

    Examples
    --------
    Create the ThermodynamicState.

    >>> from openmmtools import testsystems, states
    >>> system_container = testsystems.LysozymeImplicit()
    >>> system, positions = system_container.system, system_container.positions
    >>> thermodynamic_state = states.ThermodynamicState(system, 300*unit.kelvin)
    >>> sampler_state = states.SamplerState(positions)

    Identify ligand atoms. Topography automatically identify receptor atoms too.

    >>> from yank.yank import Topography
    >>> topography = Topography(system_container.topology, ligand_atoms=range(2603, 2621))

    you can create a completely defined restraint

    >>> restraint = Harmonic(spring_constant=8*unit.kilojoule_per_mole/unit.nanometers**2,
    ...                      restrained_receptor_atoms=[1644, 1650, 1678],
    ...                      restrained_ligand_atoms='resname TMP')

    Or automatically identify the parameters. When trying to impose a restraint
    with undefined parameters, RestraintParameterError is raised.

    >>> restraint = Harmonic()
    >>> try:
    ...     restraint.restrain_state(thermodynamic_state)
    ... except RestraintParameterError:
    ...     print('There are undefined parameters. Choosing restraint parameters automatically.')
    ...     restraint.determine_missing_parameters(thermodynamic_state, sampler_state, topography)
    ...     restraint.restrain_state(thermodynamic_state)
    ...
    There are undefined parameters. Choosing restraint parameters automatically.

    Get standard state correction.

    >>> correction = restraint.get_standard_state_correction(thermodynamic_state)

    """
    def __init__(self, spring_constant=None, **kwargs):
        super(Harmonic, self).__init__(**kwargs)
        self.spring_constant = spring_constant

    @property
    def _are_restraint_parameters_defined(self):
        """bool: True if the restraint parameters are defined."""
        return self.spring_constant is not None

    def _create_restraint_force(self, particles1, particles2):
        """Create a new restraint force between specified atoms.

        Parameters
        ----------
        particles1 : list of int
            Indices of first group of atoms to restraint.
        particles2 : list of int
            Indices of second group of atoms to restraint.

        Returns
        -------
        force : simtk.openmm.Force
           The created restraint force.

        """
        # Create bond force and lambda_restraints parameter to control it.
        if len(particles1) == 1 and len(particles2) == 1:
            # CustomCentroidBondForce works only on 64bit platforms. When the
            # restrained groups only have 1 particle, we can use the standard
            # CustomBondForce so that we can support 32bit platforms too.
            return mmtools.forces.HarmonicRestraintBondForce(spring_constant=self.spring_constant,
                                                             restrained_atom_index1=particles1[0],
                                                             restrained_atom_index2=particles2[0])
        return mmtools.forces.HarmonicRestraintForce(spring_constant=self.spring_constant,
                                                     restrained_atom_indices1=particles1,
                                                     restrained_atom_indices2=particles2)

    def _determine_restraint_parameters(self, thermodynamic_state, sampler_state, topography):
        """Automatically choose a spring constant for the restraint force.

        The spring constant is selected to give 1 kT at one standard deviation
        of receptor atoms about the receptor restrained atom.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state.
        sampler_state : openmmtools.states.SamplerState
            The sampler state holding the positions of all atoms.
        topography : yank.Topography
            The topography with labeled receptor and ligand atoms.

        """
        # Do not overwrite parameters that are already defined.
        if self.spring_constant is not None:
            return

        receptor_positions = sampler_state.positions[topography.receptor_atoms]
        sigma = pipeline.compute_radius_of_gyration(receptor_positions)

        # Compute corresponding spring constant.
        self.spring_constant = thermodynamic_state.kT / sigma**2

        logger.debug('Spring constant sigma, s = {:.3f} nm'.format(sigma / unit.nanometers))
        logger.debug('K = {:.1f} kcal/mol/A^2'.format(
            self.spring_constant / unit.kilocalories_per_mole * unit.angstroms**2))


# ==============================================================================
# Flat-bottom protein-ligand restraint.
# ==============================================================================

class FlatBottom(RadiallySymmetricRestraint):
    """A receptor-ligand restraint using a flat potential well with harmonic walls.

    An alternative choice to receptor-ligand restraints that uses a flat
    potential inside most of the protein volume with harmonic restraining
    walls outside of this. It can be used to prevent the ligand from
    drifting too far from protein in implicit solvent calculations while
    still exploring the surface of the protein for putative binding sites.

    The restraint is applied between the centroids of two groups of atoms
    that belong to the receptor and the ligand respectively. The centroids
    are determined by a mass-weighted average of the group particles positions.

    More precisely, the energy expression of the restraint is given by

        ``E = lambda_restraints * step(r-r0) * (K/2)*(r-r0)^2``

    where ``K`` is the spring constant, ``r`` is the distance between the
    restrained atoms, ``r0`` is another parameter defining the distance
    at which the restraint is imposed, and ``lambda_restraints``
    is a scale factor that can be used to control the strength of the
    restraint. You can control ``lambda_restraints`` through the class
    :class:`RestraintState`.

    The class supports automatic determination of the parameters left undefined
    in the constructor through :func:`determine_missing_parameters`.

    With OpenCL, groups with more than 1 atom are supported only on 64bit
    platforms.

    Parameters
    ----------
    spring_constant : simtk.unit.Quantity, optional
        The spring constant K (see energy expression above) in units compatible
        with joule/nanometer**2/mole (default is None).
    well_radius : simtk.unit.Quantity, optional
        The distance r0 (see energy expression above) at which the harmonic
        restraint is imposed in units of distance (default is None).
    restrained_receptor_atoms : iterable of int, int, or str, optional
        The indices of the receptor atoms to restrain, an MDTraj DSL expression, or a
        :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Select String <yank.Topography.select>`.
        This can temporarily be left undefined, but ``determine_missing_parameters()``
        must be called before using the Restraint object. The same if a DSL
        expression or Topography region is provided (default is None).
    restrained_ligand_atoms : iterable of int, int, or str, optional
        The indices of the ligand atoms to restrain, an MDTraj DSL expression.
        or a :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Select String <yank.Topography.select>`.
        This can temporarily be left undefined, but ``determine_missing_parameters()``
        must be called before using the Restraint object. The same if a DSL
        expression or Topography region is provided (default is None).

    Attributes
    ----------
    restrained_receptor_atoms : list of int or None
        The indices of the receptor atoms to restrain, an MDTraj selection string, or a Topography region selection
        string.
    restrained_ligand_atoms : list of int or None
        The indices of the ligand atoms to restrain, an MDTraj selection string, or a Topography region selection
        string.

    Examples
    --------
    Create the ThermodynamicState.

    >>> from openmmtools import testsystems, states
    >>> system_container = testsystems.LysozymeImplicit()
    >>> system, positions = system_container.system, system_container.positions
    >>> thermodynamic_state = states.ThermodynamicState(system, 298*unit.kelvin)
    >>> sampler_state = states.SamplerState(positions)

    Identify ligand atoms. Topography automatically identify receptor atoms too.

    >>> from yank.yank import Topography
    >>> topography = Topography(system_container.topology, ligand_atoms=range(2603, 2621))

    You can create a completely defined restraint

    >>> restraint = FlatBottom(spring_constant=0.6*unit.kilocalorie_per_mole/unit.angstroms**2,
    ...                        well_radius=5.2*unit.nanometers, restrained_receptor_atoms=[1644, 1650, 1678],
    ...                        restrained_ligand_atoms='resname TMP')

    or automatically identify the parameters. When trying to impose a restraint
    with undefined parameters, RestraintParameterError is raised.

    >>> restraint = FlatBottom()
    >>> try:
    ...     restraint.restrain_state(thermodynamic_state)
    ... except RestraintParameterError:
    ...     print('There are undefined parameters. Choosing restraint parameters automatically.')
    ...     restraint.determine_missing_parameters(thermodynamic_state, sampler_state, topography)
    ...     restraint.restrain_state(thermodynamic_state)
    ...
    There are undefined parameters. Choosing restraint parameters automatically.

    Get standard state correction.

    >>> correction = restraint.get_standard_state_correction(thermodynamic_state)

    """
    def __init__(self, spring_constant=None, well_radius=None, **kwargs):
        super(FlatBottom, self).__init__(**kwargs)
        self.spring_constant = spring_constant
        self.well_radius = well_radius

    @property
    def _energy_function(self):
        """str: energy expression of the restraint force."""
        return 'step(distance(g1,g2)-r0) * (K/2)*(distance(g1,g2)-r0)^2'

    @property
    def _are_restraint_parameters_defined(self):
        """bool: True if the restraint parameters are defined."""
        return self.spring_constant is not None and self.well_radius is not None

    def _create_restraint_force(self, particles1, particles2):
        """Create a new restraint force between specified atoms.

        Parameters
        ----------
        particles1 : list of int
            Indices of first group of atoms to restraint.
        particles2 : list of int
            Indices of second group of atoms to restraint.

        Returns
        -------
        force : simtk.openmm.Force
           The created restraint force.

        """
        # Create bond force and lambda_restraints parameter to control it.
        if len(particles1) == 1 and len(particles2) == 1:
            # CustomCentroidBondForce works only on 64bit platforms. When the
            # restrained groups only have 1 particle, we can use the standard
            # CustomBondForce so that we can support 32bit platforms too.
            return mmtools.forces.FlatBottomRestraintBondForce(spring_constant=self.spring_constant,
                                                               well_radius=self.well_radius,
                                                               restrained_atom_index1=particles1[0],
                                                               restrained_atom_index2=particles2[0])
        return mmtools.forces.FlatBottomRestraintForce(spring_constant=self.spring_constant,
                                                       well_radius=self.well_radius,
                                                       restrained_atom_indices1=particles1,
                                                       restrained_atom_indices2=particles2)

    def _determine_restraint_parameters(self, thermodynamic_state, sampler_state, topography):
        """Automatically choose a spring constant and well radius.

        The spring constant, is set to 5.92 kcal/mol/A**2, the well
        radius is set at twice the robust estimate of the standard
        deviation (from mean absolute deviation) plus 5 A.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state.
        sampler_state : openmmtools.states.SamplerState
            The sampler state holding the positions of all atoms.
        topography : yank.Topography
            The topography with labeled receptor and ligand atoms.

        """
        # Determine number of atoms.
        n_atoms = len(topography.receptor_atoms)

        # Check that restrained receptor atoms are in expected range.
        if any(atom_id >= n_atoms for atom_id in self.restrained_receptor_atoms):
            raise ValueError('Receptor atoms {} were selected for restraint, but system '
                             'only has {} atoms.'.format(self.restrained_receptor_atoms, n_atoms))

        # Compute well radius if the user hasn't specified it in the constructor.
        if self.well_radius is None:
            # Get positions of mass-weighted centroid atom.
            # (Working in non-unit-bearing floats for speed.)
            x_unit = sampler_state.positions.unit
            x_restrained_atoms = sampler_state.positions[self.restrained_receptor_atoms, :] / x_unit
            system = thermodynamic_state.system
            masses = np.array([system.getParticleMass(i) / unit.dalton for i in self.restrained_receptor_atoms])
            x_centroid = np.average(x_restrained_atoms, axis=0, weights=masses)

            # Get dimensionless receptor and ligand positions.
            x_receptor = sampler_state.positions[topography.receptor_atoms, :] / x_unit
            x_ligand = sampler_state.positions[topography.ligand_atoms, :] / x_unit

            # Compute maximum square distance from the centroid to any receptor atom.
            # dist2_centroid_receptor[i] is the squared distance from the centroid to receptor atom i.
            dist2_centroid_receptor = pipeline.compute_squared_distances([x_centroid], x_receptor)
            max_dist_receptor = np.sqrt(dist2_centroid_receptor.max()) * x_unit

            # Compute maximum length of the ligand. dist2_ligand_ligand[i][j] is the
            # squared distance between atoms i and j of the ligand.
            dist2_ligand_ligand = pipeline.compute_squared_distances(x_ligand, x_ligand)
            max_length_ligand = np.sqrt(dist2_ligand_ligand.max()) * x_unit

            # Compute the radius of the flat bottom restraint.
            self.well_radius = max_dist_receptor + max_length_ligand/2 + 5*unit.angstrom

        # Set default spring constant if the user hasn't specified it in the constructor.
        if self.spring_constant is None:
            self.spring_constant = 10.0 * thermodynamic_state.kT / unit.angstroms**2

        logger.debug('restraint distance r0 = {:.1f} A'.format(self.well_radius / unit.angstroms))
        logger.debug('K = {:.1f} kcal/mol/A^2'.format(
            self.spring_constant / unit.kilocalories_per_mole * unit.angstroms**2))


# ==============================================================================
# Orientation-dependent receptor-ligand restraints.
# ==============================================================================

class BoreschLike(ReceptorLigandRestraint, ABC):
    """
    Abstract class to impose Boresch-like orientational restraints on protein-ligand system.
    Subclasses are specific implementations of the energy functions

    This restraints the ligand binding mode by constraining 1 distance, 2
    angles and 3 dihedrals between 3 atoms of the receptor and 3 atoms of
    the ligand.

    More precisely, the energy expression of the restraint is given by

        .. code-block:: python

            E = lambda_restraints * {
                    K_r/2 * [|r3 - l1| - r_aA0]^2 +
                    + K_thetaA/2 * [angle(r2,r3,l1) - theta_A0]^2 +
                    + K_thetaB/2 * [angle(r3,l1,l2) - theta_B0]^2 +
                    + K_phiA/2 * hav(dihedral(r1,r2,r3,l1) - phi_A0) * 2 +
                    + K_phiB/2 * hav(dihedral(r2,r3,l1,l2) - phi_B0) * 2 +
                    + K_phiC/2 * hav(dihedral(r3,l1,l2,l3) - phi_C0) * 2
                }

    , where ``hav`` is the Haversine function ``(1-cos(x))/2`` and the parameters are:

        ``r1``, ``r2``, ``r3``: the coordinates of the 3 receptor atoms.

        ``l1``, ``l2``, ``l3``: the coordinates of the 3 ligand atoms.

        ``K_r``: the spring constant for the restrained distance ``|r3 - l1|``.

        ``r_aA0``: the equilibrium distance of ``|r3 - l1|``.

        ``K_thetaA``, ``K_thetaB``: the spring constants for ``angle(r2,r3,l1)`` and ``angle(r3,l1,l2)``.

        ``theta_A0``, ``theta_B0``: the equilibrium angles of ``angle(r2,r3,l1)`` and ``angle(r3,l1,l2)``.

        ``K_phiA``, ``K_phiB``, ``K_phiC``: the spring constants for ``dihedral(r1,r2,r3,l1)``,
        ``dihedral(r2,r3,l1,l2)``, ``dihedral(r3,l1,l2,l3)``.

        ``phi_A0``, ``phi_B0``, ``phi_C0``: the equilibrium torsion of ``dihedral(r1,r2,r3,l1)``,
        ``dihedral(r2,r3,l1,l2)``, ``dihedral(r3,l1,l2,l3)``.

        ``lambda_restraints``: a scale factor that can be used to control the strength
        of the restraint.

    You can control ``lambda_restraints`` through the class :class:`RestraintState`.

    The class supports automatic determination of the parameters left undefined
    in the constructor through :func:`determine_missing_parameters`.


    This function used to be based on the Boresch orientational restraints [1] and has similar form to
    its energy equation

        .. code-block:: python

            E = lambda_restraints * {
                    K_r/2 * [|r3 - l1| - r_aA0]^2 +
                    + K_thetaA/2 * [angle(r2,r3,l1) - theta_A0]^2 +
                    + K_thetaB/2 * [angle(r3,l1,l2) - theta_B0]^2 +
                    + K_phiA/2 * [dihedral(r1,r2,r3,l1) - phi_A0]^2 +
                    + K_phiB/2 * [dihedral(r2,r3,l1,l2) - phi_B0]^2 +
                    + K_phiC/2 * [dihedral(r3,l1,l2,l3) - phi_C0]^2
                }

    However, the form at the top is periodic with the dihedral angle and imposes a more steep energy penalty
    while still maintaining the same Taylor series expanded force and energy near phi_X0. The ``*2`` on the
    ``hav()`` functions in the energy expression are shown as the explicit correction to the ``hav()`` function
    to make the leading spring constant force consistent with the original harmonic Boresch restraint. In practice,
    the ``1/2`` from the ``hav()`` function is omitted.


    *Warning*: Symmetry corrections for symmetric ligands are not automatically applied.
    See Ref [1] and [2] for more information on correcting for ligand symmetry.

    *Warning*: Only heavy atoms can be restrained. Hydrogens will automatically be excluded.

    Parameters
    ----------
    restrained_receptor_atoms : iterable of int, str, or None; Optional
        The indices of the receptor atoms to restrain, an MDTraj DSL expression, or a
        :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Select String <yank.Topography.select>`.
        If this is a list of three ints, the receptor atoms will be restrained in order, r1, r2, r3. If there are more
        than three entries or the selection string resolves more than three atoms, the three restrained atoms will
        be chosen at random from the selection.
        This can temporarily be left undefined, but ``determine_missing_parameters()``
        must be called before using the Restraint object. The same if a DSL
        expression or Topography region is provided (default is None).
    restrained_ligand_atoms : iterable of int, str, or None; Optional
        The indices of the ligand atoms to restrain, an MDTraj DSL expression, or a
        :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Select String <yank.Topography.select>`.
        If this is a list of three ints, the receptor atoms will be restrained in order, l1, l2, l3. If there are more
        than three entries or the selection string resolves more than three atoms, the three restrained atoms will
        be chosen at random from the selection.
        This can temporarily be left undefined, but ``determine_missing_parameters()``
        must be called before using the Restraint object. The same if a DSL
        expression or Topography region is provided (default is None).
    K_r : simtk.unit.Quantity, optional
        The spring constant for the restrained distance ``|r3 - l1|`` (units
        compatible with kilocalories_per_mole/angstrom**2).
    r_aA0 : simtk.unit.Quantity, optional
        The equilibrium distance between r3 and l1 (units of length).
    K_thetaA, K_thetaB : simtk.unit.Quantity, optional
        The spring constants for ``angle(r2, r3, l1)`` and ``angle(r3, l1, l2)``
        (units compatible with kilocalories_per_mole/radians**2).
    theta_A0, theta_B0 : simtk.unit.Quantity, optional
        The equilibrium angles of ``angle(r2, r3, l1)`` and ``angle(r3, l1, l2)``
        (units compatible with radians).
    K_phiA, K_phiB, K_phiC : simtk.unit.Quantity, optional
        The spring constants for ``dihedral(r1, r2, r3, l1)``,
        ``dihedral(r2, r3, l1, l2)`` and ``dihedral(r3,l1,l2,l3)`` (units compatible
        with kilocalories_per_mole/radians**2).
    phi_A0, phi_B0, phi_C0 : simtk.unit.Quantity, optional
        The equilibrium torsion of ``dihedral(r1,r2,r3,l1)``, ``dihedral(r2,r3,l1,l2)``
        and ``dihedral(r3,l1,l2,l3)`` (units compatible with radians).

    Attributes
    ----------
    restrained_receptor_atoms : list of int
        The indices of the 3 receptor atoms to restrain [r1, r2, r3].
    restrained_ligand_atoms : list of int
        The indices of the 3 ligand atoms to restrain [l1, l2, l3].

    References
    ----------
    [1] Boresch S, Tettinger F, Leitgeb M, Karplus M. J Phys Chem B. 107:9535, 2003.
        http://dx.doi.org/10.1021/jp0217839
    [2] Mobley DL, Chodera JD, and Dill KA. J Chem Phys 125:084902, 2006.
        https://dx.doi.org/10.1063%2F1.2221683

    Examples
    --------
    Create the ThermodynamicState.

    >>> from openmmtools import testsystems, states
    >>> system_container = testsystems.LysozymeImplicit()
    >>> system, positions = system_container.system, system_container.positions
    >>> thermodynamic_state = states.ThermodynamicState(system, 298*unit.kelvin)
    >>> sampler_state = states.SamplerState(positions)

    Identify ligand atoms. Topography automatically identify receptor atoms too.

    >>> from yank.yank import Topography
    >>> topography = Topography(system_container.topology, ligand_atoms=range(2603, 2621))

    Create a partially defined restraint

    >>> restraint = Boresch(restrained_receptor_atoms=[1335, 1339, 1397],
    ...                     restrained_ligand_atoms=[2609, 2607, 2606],
    ...                     K_r=20.0*unit.kilocalories_per_mole/unit.angstrom**2,
    ...                     r_aA0=0.35*unit.nanometer)

    and automatically identify the other parameters. When trying to impose
    a restraint with undefined parameters, RestraintParameterError is raised.

    >>> try:
    ...     restraint.restrain_state(thermodynamic_state)
    ... except RestraintParameterError:
    ...     print('There are undefined parameters. Choosing restraint parameters automatically.')
    ...     restraint.determine_missing_parameters(thermodynamic_state, sampler_state, topography)
    ...     restraint.restrain_state(thermodynamic_state)
    ...
    There are undefined parameters. Choosing restraint parameters automatically.

    Get standard state correction.

    >>> correction = restraint.get_standard_state_correction(thermodynamic_state)

    """
    def __init__(self, restrained_receptor_atoms=None, restrained_ligand_atoms=None,
                 K_r=None, r_aA0=None,
                 K_thetaA=None, theta_A0=None,
                 K_thetaB=None, theta_B0=None,
                 K_phiA=None, phi_A0=None,
                 K_phiB=None, phi_B0=None,
                 K_phiC=None, phi_C0=None):
        self.restrained_receptor_atoms = restrained_receptor_atoms
        self.restrained_ligand_atoms = restrained_ligand_atoms
        self.K_r = K_r
        self.r_aA0 = r_aA0
        self.K_thetaA, self.K_thetaB = K_thetaA, K_thetaB
        self.theta_A0, self.theta_B0 = theta_A0, theta_B0
        self.K_phiA, self.K_phiB, self.K_phiC = K_phiA, K_phiB, K_phiC
        self.phi_A0, self.phi_B0, self.phi_C0 = phi_A0, phi_B0, phi_C0

    # -------------------------------------------------------------------------
    # Public properties.
    # -------------------------------------------------------------------------

    class _BoreschRestrainedAtomsProperty(_RestrainedAtomsProperty):
        """
        Descriptor of restrained atoms.

        Extends `_RestrainedAtomsProperty` to handle single integers and strings.
        """

        _MUST_COMPUTE_STRING = ('You are specifying {} {} atoms, '
                                'the final atoms will be chosen at from this set but you MUST '
                                'run "determine_missing_parameters"')

        @methoddispatch
        def _validate_atoms(self, restrained_atoms):
            restrained_atoms = super()._validate_atoms(restrained_atoms)
            if len(restrained_atoms) < 3:
                raise ValueError('At least three {} atoms are required to impose a '
                                 'Boresch-style restraint.'.format(self._atoms_type))
            elif len(restrained_atoms) > 3:
                logger.warning(self._MUST_COMPUTE_STRING.format("more than three", self._atoms_type))
            return restrained_atoms

        @_validate_atoms.register(str)
        def _cast_atom_string(self, restrained_atoms):
            logger.warning(self._MUST_COMPUTE_STRING.format("a string for", self._atoms_type))
            return restrained_atoms

    restrained_receptor_atoms = _BoreschRestrainedAtomsProperty('receptor')
    restrained_ligand_atoms = _BoreschRestrainedAtomsProperty('ligand')

    # -------------------------------------------------------------------------
    # Public methods.
    # -------------------------------------------------------------------------

    def restrain_state(self, thermodynamic_state):
        """Add the restraint force to the state's ``System``.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state holding the system to modify.

        """
        # TODO replace dihedral restraints with negative log von Mises distribution?
        #       https://en.wikipedia.org/wiki/Von_Mises_distribution, the von Mises parameter
        #       kappa would be computed from the desired standard deviation (kappa ~ sigma**(-2))
        #       and the standard state correction would need to be modified.

        # Check if all parameters are defined.
        self._check_parameters_defined()

        energy_function = self._get_energy_function_string()

        # Add constant definitions to the energy function
        for name, value in self._parameters.items():
            energy_function += '%s = %f; ' % (name, value.value_in_unit_system(unit.md_unit_system))

        # Create the force
        n_particles = 6  # number of particles involved in restraint: p1 ... p6
        restraint_force = openmm.CustomCompoundBondForce(n_particles, energy_function)
        restraint_force.addGlobalParameter('lambda_restraints', 1.0)
        restraint_force.addBond(self.restrained_receptor_atoms + self.restrained_ligand_atoms, [])
        restraint_force.setUsesPeriodicBoundaryConditions(thermodynamic_state.is_periodic)

        # Get a copy of the system of the ThermodynamicState, modify it and set it back.
        system = thermodynamic_state.system
        self._add_force_in_separate_group(system, restraint_force)
        thermodynamic_state.system = system

    def get_standard_state_correction(self, thermodynamic_state):
        """Return the standard state correction.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state.

        Returns
        -------
        DeltaG : float
           Computed standard-state correction in dimensionless units (kT).

        """
        self._check_parameters_defined()

        def strip(passed_unit):
            """Cast the passed_unit into md unit system for integrand lambda functions"""
            return passed_unit.value_in_unit_system(unit.md_unit_system)

        # Shortcuts variables.
        pi = np.pi
        kT = thermodynamic_state.kT
        p = self  # For the parameters.

        # Radial
        sigma = 1 / unit.sqrt(p.K_r / kT)
        rmin = strip(max(0*unit.angstrom, p.r_aA0 - 8 * sigma))
        rmax = strip(p.r_aA0 + 8 * sigma)
        r0 = strip(p.r_aA0)
        K_r = strip(p.K_r)
        I = lambda r: self._numerical_distance_integrand(r, r0, K_r, strip(kT))
        integral_packet = scipy.integrate.quad(I, rmin, rmax) * unit.nanometer**3
        ExpDeltaG = integral_packet[0]

        # Angular
        for name in ['A', 'B']:
            theta0 = strip(getattr(p, 'theta_' + name + '0'))
            K_theta = strip(getattr(p, 'K_theta' + name))
            I = lambda theta: self._numerical_angle_integrand(theta, theta0, K_theta, strip(kT))
            integral_packet = scipy.integrate.quad(I, 0, pi)
            ExpDeltaG *= integral_packet[0]

        # Torsion
        for name in ['A', 'B', 'C']:
            phi0 = strip(getattr(p, 'phi_' + name + '0'))
            K_phi = strip(getattr(p, 'K_phi' + name))
            I = lambda phi: self._numerical_torsion_integrand(phi, phi0, K_phi, strip(kT))
            integral_packet = scipy.integrate.quad(I, -pi, pi)
            ExpDeltaG *= integral_packet[0]

        DeltaG = -np.log(8 * pi**2 * V0 / ExpDeltaG)
        return DeltaG

    def determine_missing_parameters(self, thermodynamic_state, sampler_state, topography):
        """Determine parameters and restrained atoms automatically.

        Currently, all equilibrium values are measured from the initial structure,
        while spring constants are set to 20 kcal/(mol A**2) or 20 kcal/(mol rad**2)
        as in Ref [1]. The restrained atoms are selected so that the analytical
        standard state correction will be valid.

        Parameters that have been already specified are left untouched.

        Future iterations of this feature will introduce the ability to extract
        equilibrium parameters and spring constants from a short simulation.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state.
        sampler_state : openmmtools.states.SamplerState, optional
            The sampler state holding the positions of all atoms.
        topography : yank.Topography, optional
            The topography with labeled receptor and ligand atoms.

        """

        logger.debug('Automatically selecting restraint atoms and parameters:')

        # If restrained atoms are already specified, we only need to determine parameters.
        if self._are_restrained_atoms_defined:
            self._determine_restraint_parameters(sampler_state, topography)
        else:
            restrained_atoms = self._pick_restrained_atoms(sampler_state, topography)
            self.restrained_receptor_atoms = restrained_atoms[:3]
            self.restrained_ligand_atoms = restrained_atoms[3:]

            # Determine restraint parameters for these atoms.
            self._determine_restraint_parameters(sampler_state, topography)

        ssc = self.get_standard_state_correction(thermodynamic_state)
        logger.debug('Standard state correction: {} kT'.format(ssc))

    # -------------------------------------------------------------------------
    # Abstract Functions
    # -------------------------------------------------------------------------

    @abc.abstractmethod
    def _get_energy_function_string(self):
        """
        Get the energy function string which defines the full restraint compatible with OpenMM Custom*Force
        expressions.

        Restraint constants can be used in this function and will be substituted in

          * K_r, r_aA0
          * K_thetaA, theta_A0,
          * K_thetaB, theta_B0,
          * K_phiA, phi_A0,
          * K_phiB, phi_B0,
          * K_phiC, phi_C0

        Variables should also be used:

          * lambda_restraints    : Alchemical variable, should be scalar on whole energy
          * distance(p3,p4)      : Restrained distance
          * angle(p2,p3,p4)      : Restrained angle "A"
          * angle(p3,p4,p5)      : Restrained angle "B"
          * dihedral(p1,p2,p3,p4): Restrained torsion "A"
          * dihedral(p2,p3,p4,p5): Restrained torsion "B"
          * dihedral(p3,p4,p5,p6): Restrained torsion "C"

        Returns
        -------
        energy_function : string
            String defining the force compatible with OpenMM Custom
        """
        pass

    @abc.abstractmethod
    def _numerical_torsion_integrand(self, phi, phi0, spring_constant, kt):
        """
        Integrand for the torsion (phi) restraints which will be integrated numerically for standard state correction

        Domain is on [-pi, pi], the same domain OpenMM uses

        Parameters
        ----------
        phi : float or np.ndarray of float
            Torsion angle which will be integrated, units of radians
        phi0 : float
            Equilibrium torsion angle at which force of restraint often is 0, units of radians
        spring_constant : float
            Spring constant for this torsion in units of kJ/mol/nm**2
        kt : float
            Boltzmann Temperature of the thermodynamic state restraining the atoms = kB * T
            in units of kJ/mol

        Returns
        -------
        integrand : float
            Value of the integrated
        """
        pass

    @abc.abstractmethod
    def _numerical_angle_integrand(self, theta, theta0, spring_constant, kt):
        """
        Integrand for the angle (theta) restraints which will be integrated numerically for standard state correction

        Domain is on [0, pi]

        Parameters
        ----------
        theta : float or np.ndarray of float
            Angle which will be integrated, units of radians
        theta0 : float
            Equilibrium angle at which force of restraint often is 0, units of radians
        spring_constant : float
            Spring constant for this angle in units of with kJ/mol/nm**2
        kt : float
            Boltzmann Temperature of the thermodynamic state restraining the atoms = kB * T
            in units of kJ/mol

        Returns
        -------
        integrand : float
            Value of the integrated
        """
        pass

    @abc.abstractmethod
    def _numerical_distance_integrand(self, r, r0, spring_constant, kt):
        """
        Integrand for the distance restraint which will be integrated numerically for standard state correction

        Domain is on [0, +infinity]

        Parameters
        ----------
        r : float or np.ndarray of float
           Distance which will be integrated, units of nm
        r0 : float
            Equilibrium distance at which force of restraint often is 0, nm
        spring_constant : float
            Spring constant for this distance in units of kJ/mol/nm**2
        kt : float
            Boltzmann Temperature of the thermodynamic state restraining the atoms = kB * T
            in units of kJ/mol

        Returns
        -------
        integrand : float
            Value of the integrated
        """
        pass

    # -------------------------------------------------------------------------
    # Internal-usage
    # -------------------------------------------------------------------------

    def _check_parameters_defined(self):
        """Raise an exception there are still parameters undefined."""
        if not self._are_restrained_atoms_defined:
            raise RestraintParameterError('Undefined restrained atoms.')

        # Find undefined parameters and raise error.
        undefined_parameters = [name for name, value in self._parameters.items() if value is None]
        if len(undefined_parameters) > 0:
            err_msg = 'Undefined parameters for Boresch-like restraint: {}'.format(undefined_parameters)
            raise RestraintParameterError(err_msg)

    @property
    def _are_restrained_atoms_defined(self):
        """Check if the restrained atoms are defined well enough to make a restraint"""
        for atoms in [self.restrained_receptor_atoms, self.restrained_ligand_atoms]:
            # Atoms should be a list or None at this point due to the _RestrainedAtomsProperty class
            if atoms is None or not (isinstance(atoms, list) and len(atoms) == 3):
                return False
        return True

    @staticmethod
    def _is_collinear(positions, atoms, threshold=0.9):
        """Report whether any sequential vectors in a sequence of atoms are collinear.

        Parameters
        ----------
        positions : n_atoms x 3 simtk.unit.Quantity
            Reference positions to use for imposing restraints (units of length).
        atoms : iterable of int
            The indices of the atoms to test.
        threshold : float, optional, default=0.9
            Atoms are not collinear if their sequential vector separation dot
            products are less than ``threshold``.

        Returns
        -------
        result : bool
            Returns True if any sequential pair of vectors is collinear; False otherwise.

        """
        result = False
        for i in range(len(atoms)-2):
            v1 = positions[atoms[i+1], :] - positions[atoms[i], :]
            v2 = positions[atoms[i+2], :] - positions[atoms[i+1], :]
            normalized_inner_product = np.dot(v1, v2) / np.sqrt(np.dot(v1, v1) * np.dot(v2, v2))
            result = result or (normalized_inner_product > threshold)

        return result

    def _pick_restrained_atoms(self, sampler_state, topography):
        """Select atoms to be used in restraint.

        Parameters
        ----------
        sampler_state : openmmtools.states.SamplerState, optional
            The sampler state holding the positions of all atoms.
        topography : yank.Topography, optional
            The topography with labeled receptor and ligand atoms.

        Returns
        -------
        restrained_atoms : list of int
            List of six atom indices used in the restraint.
            restrained_atoms[0:3] belong to the receptor,
            restrained_atoms[4:6] belong to the ligand.

        Notes
        -----
        The current algorithm simply selects random subsets of receptor
        and ligand atoms and rejects those that are too close to collinear.
        Future updates can further refine this algorithm.

        """
        # If receptor and ligand atoms are explicitly provided, use those.
        heavy_ligand_atoms = self.restrained_ligand_atoms
        heavy_receptor_atoms = self.restrained_receptor_atoms

        # Otherwise we restrain only heavy atoms.
        heavy_atoms = set(topography.topology.select('not element H').tolist())
        # Intersect heavy atoms with receptor/ligand atoms (s1&s2 is intersect).

        atom_selector = _AtomSelector(topography)

        heavy_ligand_atoms = atom_selector.compute_atom_intersect(heavy_ligand_atoms, 'ligand_atoms', heavy_atoms)
        heavy_receptor_atoms = atom_selector.compute_atom_intersect(heavy_receptor_atoms, 'receptor_atoms', heavy_atoms)

        if len(heavy_receptor_atoms) < 3 or len(heavy_ligand_atoms) < 3:
            raise ValueError('There must be at least three heavy atoms in receptor_atoms '
                             '(# heavy {}) and ligand_atoms (# heavy {}).'.format(
                                     len(heavy_receptor_atoms), len(heavy_ligand_atoms)))

        # If r3 or l1 atoms are given. We have to pick those.
        if isinstance(heavy_receptor_atoms, list):
            r3_atoms = [heavy_receptor_atoms[2]]
        else:
            r3_atoms = heavy_receptor_atoms
        if isinstance(heavy_ligand_atoms, list):
            l1_atoms = [heavy_ligand_atoms[0]]
        else:
            l1_atoms = heavy_ligand_atoms
        # TODO: Cast itertools generator to np array more efficiently
        r3_l1_pairs = np.array(list(itertools.product(r3_atoms, l1_atoms)))

        # Filter r3-l1 pairs that are too close/far away for the distance constraint.
        max_distance = 4 * unit.angstrom/unit.nanometer
        min_distance = 1 * unit.angstrom/unit.nanometer
        t = md.Trajectory(sampler_state.positions / unit.nanometers, topography.topology)
        distances = md.geometry.compute_distances(t, r3_l1_pairs)[0]
        indices_of_in_range_pairs = np.where(np.logical_and(distances > min_distance, distances <= max_distance))[0]

        if len(indices_of_in_range_pairs) == 0:
            error_msg = ('There are no heavy ligand atoms within the range of [{},{}] nm heavy receptor atoms!\n'
                         'Please Check your input files or try another restraint class')
            raise ValueError(error_msg.format(min_distance, max_distance))
        r3_l1_pairs = r3_l1_pairs[indices_of_in_range_pairs].tolist()

        def find_bonded_to(input_atom_index, comparison_set):
            """
            Find bonded network between the atoms to create a selection with 1 angle to the reference

            Parameters
            ----------
            input_atom_index : int
                Reference atom index to try and create the selection from the bonds
            comparison_set : iterable of int
                Set of additional atoms to try and make the selection from. There should be at least
                one non-colinear set 3 atoms which are bonded together in R-B-C where R is the input_atom_index
                and B, C are atoms in the comparison_set bonded to each other.
                Can be inclusive of input_atom_index and C can be bound to R as well as B

            Returns
            -------
            bonded_atoms : list of int, length 3
                Returns the list of atoms in order of input_atom_index <- bonded atom <- bonded atom
            """
            # Probably could make this faster if we added a graph module like networkx dep, but not needed
            # Could also be done by iterating over OpenMM System angles
            # Get topology
            top = topography.topology
            bonds = np.zeros([top.n_atoms, top.n_atoms], dtype=bool)
            # Create bond graph
            for a1, a2 in top.bonds:
                a1 = a1.index
                a2 = a2.index
                bonds[a1, a2] = bonds[a2, a1] = True
            all_bond_options = []
            # Cycle through all bonds on the reference
            for a2, first_bond in enumerate(bonds[input_atom_index]):
                # Enumerate all secondary bonds from the reference but only if in comparison set
                if first_bond and a2 in comparison_set:
                    # Same as first
                    for a3, second_bond in enumerate(bonds[a2]):
                        if second_bond and a3 in comparison_set and a3 != input_atom_index:
                            all_bond_options.append([a2, a3])
            # This will raise a ValueError if nothing is found
            return random.sample(all_bond_options, 1)[0]

        # Iterate until we have found a set of non-collinear atoms.
        accepted = False
        max_attempts = 100
        attempts = 0
        while not accepted:
            logger.debug('Attempt {} / {} at automatically selecting atoms and '
                         'restraint parameters...'.format(attempts, max_attempts))
            # Select a receptor/ligand atom in range of each other for the distance constraint.
            r3_l1_atoms = random.sample(r3_l1_pairs, 1)[0]
            # Determine remaining receptor/ligand atoms.
            if isinstance(heavy_receptor_atoms, list):
                r1_r2_atoms = heavy_receptor_atoms[:2]
            else:
                try:
                    r1_r2_atoms = find_bonded_to(r3_l1_atoms[0], heavy_receptor_atoms)[::-1]
                except ValueError:
                    r1_r2_atoms = None
            if isinstance(heavy_ligand_atoms, list):
                l2_l3_atoms = heavy_ligand_atoms[1:]
            else:
                try:
                    l2_l3_atoms = find_bonded_to(r3_l1_atoms[-1], heavy_ligand_atoms)
                except ValueError:
                    l2_l3_atoms = None
            # Reject collinear sets of atoms.
            if r1_r2_atoms is None or l2_l3_atoms is None:
                accepted = False
            else:
                restrained_atoms = r1_r2_atoms + r3_l1_atoms + l2_l3_atoms
                accepted = not self._is_collinear(sampler_state.positions, restrained_atoms)
            if attempts > max_attempts:
                raise RuntimeError("Could not find any good sets of bonded atoms to make stable Boresch-like "
                                   "restraints from. There should be at least 1 real defined angle in the"
                                   "selected restrained ligand atoms and 1 in the selected restrained receptor atoms "
                                   "for good numerical stability")
            else:
                attempts += 1

        logger.debug('Selected atoms to restrain: {}'.format(restrained_atoms))
        return restrained_atoms

    def _determine_restraint_parameters(self, sampler_states, topography):
        """Determine restraint parameters.

        Currently, all equilibrium values are measured from the initial structure,
        while spring constants are set to 20 kcal/(mol A**2) or 20 kcal/(mol rad**2)
        as in [1].

        Future iterations of this feature will introduce the ability to extract
        equilibrium parameters and spring constants from a short simulation.

        References
        ----------
        [1] Boresch S, Tettinger F, Leitgeb M, Karplus M. J Phys Chem B. 107:9535, 2003.
        http://dx.doi.org/10.1021/jp0217839

        """
        # We determine automatically only the parameters that have been left undefined.
        def _assign_if_undefined(attr_name, attr_value):
            """Assign value to self.name only if it is None."""
            if getattr(self, attr_name) is None:
                setattr(self, attr_name, attr_value)

        # Merge receptor and ligand atoms in a single array for easy manipulation.
        restrained_atoms = self.restrained_receptor_atoms + self.restrained_ligand_atoms

        # Set spring constants uniformly, as in Ref [1] Table 1 caption.
        _assign_if_undefined('K_r', 20.0 * unit.kilocalories_per_mole / unit.angstrom**2)
        for parameter_name in ['K_thetaA', 'K_thetaB', 'K_phiA', 'K_phiB', 'K_phiC']:
            _assign_if_undefined(parameter_name, 20.0 * unit.kilocalories_per_mole / unit.radian**2)

        # Measure equilibrium geometries from static reference structure
        t = md.Trajectory(sampler_states.positions / unit.nanometers, topography.topology)

        atom_pairs = [restrained_atoms[2:4]]
        distances = md.geometry.compute_distances(t, atom_pairs, periodic=False)
        _assign_if_undefined('r_aA0', distances[0][0] * unit.nanometers)

        atom_triplets = [restrained_atoms[i:(i+3)] for i in range(1, 3)]
        angles = md.geometry.compute_angles(t, atom_triplets, periodic=False)
        for parameter_name, angle in zip(['theta_A0', 'theta_B0'], angles[0]):
            _assign_if_undefined(parameter_name, angle * unit.radians)

        atom_quadruplets = [restrained_atoms[i:(i+4)] for i in range(3)]
        dihedrals = md.geometry.compute_dihedrals(t, atom_quadruplets, periodic=False)
        for parameter_name, angle in zip(['phi_A0', 'phi_B0', 'phi_C0'], dihedrals[0]):
            _assign_if_undefined(parameter_name, angle * unit.radians)

        # Write restraint parameters
        msg = 'restraint parameters:\n'
        for parameter_name, parameter_value in self._parameters.items():
            msg += '%24s : %s\n' % (parameter_name, parameter_value)
        logger.debug(msg)


class Boresch(BoreschLike):
    """
    Impose Boresch-style orientational restraints on protein-ligand system.

    This restraints the ligand binding mode by constraining 1 distance, 2
    angles and 3 dihedrals between 3 atoms of the receptor and 3 atoms of
    the ligand.

    More precisely, the energy expression of the restraint is given by

        .. code-block:: python

            E = lambda_restraints * {
                    K_r/2 * [|r3 - l1| - r_aA0]^2 +
                    + K_thetaA/2 * [angle(r2,r3,l1) - theta_A0]^2 +
                    + K_thetaB/2 * [angle(r3,l1,l2) - theta_B0]^2 +
                    + K_phiA/2 * [dihedral(r1,r2,r3,l1) - phi_A0]^2 +
                    + K_phiB/2 * [dihedral(r2,r3,l1,l2) - phi_B0]^2 +
                    + K_phiC/2 * [dihedral(r3,l1,l2,l3) - phi_C0]^2
                }

    , where the parameters are:

        ``r1``, ``r2``, ``r3``: the coordinates of the 3 receptor atoms.

        ``l1``, ``l2``, ``l3``: the coordinates of the 3 ligand atoms.

        ``K_r``: the spring constant for the restrained distance ``|r3 - l1|``.

        ``r_aA0``: the equilibrium distance of ``|r3 - l1|``.

        ``K_thetaA``, ``K_thetaB``: the spring constants for ``angle(r2,r3,l1)`` and ``angle(r3,l1,l2)``.

        ``theta_A0``, ``theta_B0``: the equilibrium angles of ``angle(r2,r3,l1)`` and ``angle(r3,l1,l2)``.

        ``K_phiA``, ``K_phiB``, ``K_phiC``: the spring constants for ``dihedral(r1,r2,r3,l1)``,
        ``dihedral(r2,r3,l1,l2)``, ``dihedral(r3,l1,l2,l3)``.

        ``phi_A0``, ``phi_B0``, ``phi_C0``: the equilibrium torsion of ``dihedral(r1,r2,r3,l1)``,
        ``dihedral(r2,r3,l1,l2)``, ``dihedral(r3,l1,l2,l3)``.

        ``lambda_restraints``: a scale factor that can be used to control the strength
        of the restraint.

    You can control ``lambda_restraints`` through the class :class:`RestraintState`.

    The class supports automatic determination of the parameters left undefined
    in the constructor through :func:`determine_missing_parameters`.

    *Warning*: Symmetry corrections for symmetric ligands are not automatically applied.
    See Ref [1] and [2] for more information on correcting for ligand symmetry.

    *Warning*: Only heavy atoms can be restrained. Hydrogens will automatically be excluded.

    Parameters
    ----------
    restrained_receptor_atoms : iterable of int, str, or None; Optional
        The indices of the receptor atoms to restrain, an MDTraj DSL expression, or a
        :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Select String <yank.Topography.select>`.
        If this is a list of three ints, the receptor atoms will be restrained in order, r1, r2, r3. If there are more
        than three entries or the selection string resolves more than three atoms, the three restrained atoms will
        be chosen at random from the selection.
        This can temporarily be left undefined, but ``determine_missing_parameters()``
        must be called before using the Restraint object. The same if a DSL
        expression or Topography region is provided (default is None).
    restrained_ligand_atoms : iterable of int, str, or None; Optional
        The indices of the ligand atoms to restrain, an MDTraj DSL expression, or a
        :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Select String <yank.Topography.select>`.
        If this is a list of three ints, the receptor atoms will be restrained in order, l1, l2, l3. If there are more
        than three entries or the selection string resolves more than three atoms, the three restrained atoms will
        be chosen at random from the selection.
        This can temporarily be left undefined, but ``determine_missing_parameters()``
        must be called before using the Restraint object. The same if a DSL
        expression or Topography region is provided (default is None).
    K_r : simtk.unit.Quantity, optional
        The spring constant for the restrained distance ``|r3 - l1|`` (units
        compatible with kilocalories_per_mole/angstrom**2).
    r_aA0 : simtk.unit.Quantity, optional
        The equilibrium distance between r3 and l1 (units of length).
    K_thetaA, K_thetaB : simtk.unit.Quantity, optional
        The spring constants for ``angle(r2, r3, l1)`` and ``angle(r3, l1, l2)``
        (units compatible with kilocalories_per_mole/radians**2).
    theta_A0, theta_B0 : simtk.unit.Quantity, optional
        The equilibrium angles of ``angle(r2, r3, l1)`` and ``angle(r3, l1, l2)``
        (units compatible with radians).
    K_phiA, K_phiB, K_phiC : simtk.unit.Quantity, optional
        The spring constants for ``dihedral(r1, r2, r3, l1)``,
        ``dihedral(r2, r3, l1, l2)`` and ``dihedral(r3,l1,l2,l3)`` (units compatible
        with kilocalories_per_mole/radians**2).
    phi_A0, phi_B0, phi_C0 : simtk.unit.Quantity, optional
        The equilibrium torsion of ``dihedral(r1,r2,r3,l1)``, ``dihedral(r2,r3,l1,l2)``
        and ``dihedral(r3,l1,l2,l3)`` (units compatible with radians).

    Attributes
    ----------
    restrained_receptor_atoms : list of int
        The indices of the 3 receptor atoms to restrain [r1, r2, r3].
    restrained_ligand_atoms : list of int
        The indices of the 3 ligand atoms to restrain [l1, l2, l3].

    References
    ----------
    [1] Boresch S, Tettinger F, Leitgeb M, Karplus M. J Phys Chem B. 107:9535, 2003.
        http://dx.doi.org/10.1021/jp0217839
    [2] Mobley DL, Chodera JD, and Dill KA. J Chem Phys 125:084902, 2006.
        https://dx.doi.org/10.1063%2F1.2221683

    Examples
    --------
    Create the ThermodynamicState.

    >>> from openmmtools import testsystems, states
    >>> system_container = testsystems.LysozymeImplicit()
    >>> system, positions = system_container.system, system_container.positions
    >>> thermodynamic_state = states.ThermodynamicState(system, 298*unit.kelvin)
    >>> sampler_state = states.SamplerState(positions)

    Identify ligand atoms. Topography automatically identify receptor atoms too.

    >>> from yank.yank import Topography
    >>> topography = Topography(system_container.topology, ligand_atoms=range(2603, 2621))

    Create a partially defined restraint

    >>> restraint = Boresch(restrained_receptor_atoms=[1335, 1339, 1397],
    ...                     restrained_ligand_atoms=[2609, 2607, 2606],
    ...                     K_r=20.0*unit.kilocalories_per_mole/unit.angstrom**2,
    ...                     r_aA0=0.35*unit.nanometer)

    and automatically identify the other parameters. When trying to impose
    a restraint with undefined parameters, RestraintParameterError is raised.

    >>> try:
    ...     restraint.restrain_state(thermodynamic_state)
    ... except RestraintParameterError:
    ...     print('There are undefined parameters. Choosing restraint parameters automatically.')
    ...     restraint.determine_missing_parameters(thermodynamic_state, sampler_state, topography)
    ...     restraint.restrain_state(thermodynamic_state)
    ...
    There are undefined parameters. Choosing restraint parameters automatically.

    Get standard state correction.

    >>> correction = restraint.get_standard_state_correction(thermodynamic_state)
    """

    def _get_energy_function_string(self):
        """
        Get the energy function string which defines the full restraint compatible with OpenMM Custom*Force
        expressions.

        Returns
        -------
        energy_function : string
            String defining the force compatible with OpenMM Custom
        """
        energy_function = """
            lambda_restraints * E;
            E = (K_r/2)*(distance(p3,p4) - r_aA0)^2
            + (K_thetaA/2)*(angle(p2,p3,p4)-theta_A0)^2 + (K_thetaB/2)*(angle(p3,p4,p5)-theta_B0)^2
            + (K_phiA/2)*dphi_A^2 + (K_phiB/2)*dphi_B^2 + (K_phiC/2)*dphi_C^2;
            dphi_A = dA - floor(dA/(2*pi)+0.5)*(2*pi); dA = dihedral(p1,p2,p3,p4) - phi_A0;
            dphi_B = dB - floor(dB/(2*pi)+0.5)*(2*pi); dB = dihedral(p2,p3,p4,p5) - phi_B0;
            dphi_C = dC - floor(dC/(2*pi)+0.5)*(2*pi); dC = dihedral(p3,p4,p5,p6) - phi_C0;
            pi = %f;
            """ % np.pi
        return energy_function

    def _numerical_distance_integrand(self, r, r0, spring_constant, kt):
        """
        Integrand for the distance restraint which will be integrated numerically for standard state correction

        Domain is on [0, +infinity], practically is only taken up to 8 thermal fluctuations.

        Parameters
        ----------
        r : float or np.ndarray of float
           Distance which will be integrated, units of nm
        r0 : float
            Equilibrium distance at which force of restraint often is 0, nm
        spring_constant : float
            Spring constant for this distance in units of kJ/mol/nm**2
        kt : float
            Boltzmann Temperature of the thermodynamic state restraining the atoms = kB * T
            in units of kJ/mol

        Returns
        -------
        integrand : float
            Value of the integrated
        """
        return r ** 2 * np.exp(-spring_constant / (2 * kt) * (r - r0) ** 2)

    def _numerical_angle_integrand(self, theta, theta0, spring_constant, kt):
        """
        Integrand for the angle (theta) restraints which will be integrated numerically for standard state correction

        This uses a harmonic restraint centered around theta0

        Domain is on [0, pi]

        Parameters
        ----------
        theta : float or np.ndarray of float
            Angle which will be integrated, units of radians
        theta0 : float
            Equilibrium angle at which force of restraint often is 0, units of radians
        spring_constant : float
            Spring constant for this angle in units of kJ/mol/nm**2

        Returns
        -------
        integrand : float
            Value of the integrated
        """
        return np.sin(theta) * np.exp(-spring_constant / (2 * kt) * (theta - theta0) ** 2)

    def _numerical_torsion_integrand(self, phi, phi0, spring_constant, kt):
        """
        Integrand for the torsion (phi) restraints which will be integrated numerically for standard state correction

        Uses a harmonic restraint around phi0 with corrections to the OpenMM limits of dihedral calculation

        Domain is on [-pi, pi], the same domain OpenMM uses

        Parameters
        ----------
        phi : float or np.ndarray of float
            Torsion angle which will be integrated, units of radians
        phi0 : float
            Equilibrium torsion angle at which force of restraint often is 0, units of radians
        spring_constant : float
            Spring constant for this torsion in units of kJ/mol/nm**2
        kt : float
            Boltzmann Temperature of the thermodynamic state restraining the atoms = kB * T
            in units of kJ/mol

        Returns
        -------
        integrand : float
            Value of the integrated
        """
        d_tor = phi - phi0
        dphi = d_tor - np.floor(d_tor / (2 * np.pi) + 0.5) * (2 * np.pi)
        return np.exp(-(spring_constant / 2) / kt * dphi ** 2)


class PeriodicTorsionBoresch(Boresch):
    """
    Impose Boresch-style orientational restraints on protein-ligand system where torsions are restrained by a
    periodic instead of harmonic force

    This restraints the ligand binding mode by constraining 1 distance, 2
    angles and 3 dihedrals between 3 atoms of the receptor and 3 atoms of
    the ligand.

    More precisely, the energy expression of the restraint is given by

        .. code-block:: python

            E = lambda_restraints * {
                    K_r/2 * [|r3 - l1| - r_aA0]^2 +
                    + K_thetaA/2 * [angle(r2,r3,l1) - theta_A0]^2 +
                    + K_thetaB/2 * [angle(r3,l1,l2) - theta_B0]^2 +
                    + K_phiA/2 * hav(dihedral(r1,r2,r3,l1) - phi_A0) * 2 +
                    + K_phiB/2 * hav(dihedral(r2,r3,l1,l2) - phi_B0) * 2 +
                    + K_phiC/2 * hav(dihedral(r3,l1,l2,l3) - phi_C0) * 2
                }

    , where ``hav`` is the Haversine function ``(1-cos(x))/2`` and the parameters are:

        ``r1``, ``r2``, ``r3``: the coordinates of the 3 receptor atoms.

        ``l1``, ``l2``, ``l3``: the coordinates of the 3 ligand atoms.

        ``K_r``: the spring constant for the restrained distance ``|r3 - l1|``.

        ``r_aA0``: the equilibrium distance of ``|r3 - l1|``.

        ``K_thetaA``, ``K_thetaB``: the spring constants for ``angle(r2,r3,l1)`` and ``angle(r3,l1,l2)``.

        ``theta_A0``, ``theta_B0``: the equilibrium angles of ``angle(r2,r3,l1)`` and ``angle(r3,l1,l2)``.

        ``K_phiA``, ``K_phiB``, ``K_phiC``: the spring constants for ``dihedral(r1,r2,r3,l1)``,
        ``dihedral(r2,r3,l1,l2)``, ``dihedral(r3,l1,l2,l3)``.

        ``phi_A0``, ``phi_B0``, ``phi_C0``: the equilibrium torsion of ``dihedral(r1,r2,r3,l1)``,
        ``dihedral(r2,r3,l1,l2)``, ``dihedral(r3,l1,l2,l3)``.

        ``lambda_restraints``: a scale factor that can be used to control the strength
        of the restraint.

    You can control ``lambda_restraints`` through the class :class:`RestraintState`.

    The class supports automatic determination of the parameters left undefined
    in the constructor through :func:`determine_missing_parameters`.


    This function used to be based on the Boresch orientational restraints [1] and has similar form to
    its energy equation

        .. code-block:: python

            E = lambda_restraints * {
                    K_r/2 * [|r3 - l1| - r_aA0]^2 +
                    + K_thetaA/2 * [angle(r2,r3,l1) - theta_A0]^2 +
                    + K_thetaB/2 * [angle(r3,l1,l2) - theta_B0]^2 +
                    + K_phiA/2 * [dihedral(r1,r2,r3,l1) - phi_A0]^2 +
                    + K_phiB/2 * [dihedral(r2,r3,l1,l2) - phi_B0]^2 +
                    + K_phiC/2 * [dihedral(r3,l1,l2,l3) - phi_C0]^2
                }

    However, the form at the top is periodic with the dihedral angle and imposes a more steep energy penalty
    while still maintaining the same Taylor series expanded force and energy near phi_X0. The ``*2`` on the
    ``hav()`` functions in the energy expression are shown as the explicit correction to the ``hav()`` function
    to make the leading spring constant force consistent with the original harmonic Boresch restraint. In practice,
    the ``1/2`` from the ``hav()`` function is omitted.


    *Warning*: Symmetry corrections for symmetric ligands are not automatically applied.
    See Ref [1] and [2] for more information on correcting for ligand symmetry.

    *Warning*: Only heavy atoms can be restrained. Hydrogens will automatically be excluded.

    Parameters
    ----------
    restrained_receptor_atoms : iterable of int, str, or None; Optional
        The indices of the receptor atoms to restrain, an MDTraj DSL expression, or a
        :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Select String <yank.Topography.select>`.
        If this is a list of three ints, the receptor atoms will be restrained in order, r1, r2, r3. If there are more
        than three entries or the selection string resolves more than three atoms, the three restrained atoms will
        be chosen at random from the selection.
        This can temporarily be left undefined, but ``determine_missing_parameters()``
        must be called before using the Restraint object. The same if a DSL
        expression or Topography region is provided (default is None).
    restrained_ligand_atoms : iterable of int, str, or None; Optional
        The indices of the ligand atoms to restrain, an MDTraj DSL expression, or a
        :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Select String <yank.Topography.select>`.
        If this is a list of three ints, the receptor atoms will be restrained in order, l1, l2, l3. If there are more
        than three entries or the selection string resolves more than three atoms, the three restrained atoms will
        be chosen at random from the selection.
        This can temporarily be left undefined, but ``determine_missing_parameters()``
        must be called before using the Restraint object. The same if a DSL
        expression or Topography region is provided (default is None).
    K_r : simtk.unit.Quantity, optional
        The spring constant for the restrained distance ``|r3 - l1|`` (units
        compatible with kilocalories_per_mole/angstrom**2).
    r_aA0 : simtk.unit.Quantity, optional
        The equilibrium distance between r3 and l1 (units of length).
    K_thetaA, K_thetaB : simtk.unit.Quantity, optional
        The spring constants for ``angle(r2, r3, l1)`` and ``angle(r3, l1, l2)``
        (units compatible with kilocalories_per_mole/radians**2).
    theta_A0, theta_B0 : simtk.unit.Quantity, optional
        The equilibrium angles of ``angle(r2, r3, l1)`` and ``angle(r3, l1, l2)``
        (units compatible with radians).
    K_phiA, K_phiB, K_phiC : simtk.unit.Quantity, optional
        The spring constants for ``dihedral(r1, r2, r3, l1)``,
        ``dihedral(r2, r3, l1, l2)`` and ``dihedral(r3,l1,l2,l3)`` (units compatible
        with kilocalories_per_mole/radians**2).
    phi_A0, phi_B0, phi_C0 : simtk.unit.Quantity, optional
        The equilibrium torsion of ``dihedral(r1,r2,r3,l1)``, ``dihedral(r2,r3,l1,l2)``
        and ``dihedral(r3,l1,l2,l3)`` (units compatible with radians).

    Attributes
    ----------
    restrained_receptor_atoms : list of int
        The indices of the 3 receptor atoms to restrain [r1, r2, r3].
    restrained_ligand_atoms : list of int
        The indices of the 3 ligand atoms to restrain [l1, l2, l3].

    References
    ----------
    [1] Boresch S, Tettinger F, Leitgeb M, Karplus M. J Phys Chem B. 107:9535, 2003.
        http://dx.doi.org/10.1021/jp0217839
    [2] Mobley DL, Chodera JD, and Dill KA. J Chem Phys 125:084902, 2006.
        https://dx.doi.org/10.1063%2F1.2221683
    """

    def _get_energy_function_string(self):
        """
        Get the energy function string which defines the full restraint compatible with OpenMM Custom*Force
        expressions.

        Returns
        -------
        energy_function : string
            String defining the force compatible with OpenMM Custom
        """
        energy_function = """
            lambda_restraints * E;
            E = (K_r/2)*(distance(p3,p4) - r_aA0)^2
            + (K_thetaA/2)*(angle(p2,p3,p4)-theta_A0)^2 + (K_thetaB/2)*(angle(p3,p4,p5)-theta_B0)^2
            + (K_phiA/2)*uphi_A + (K_phiB/2)*uphi_B + (K_phiC/2)*uphi_C;
            uphi_A = (1-cos(dA)); dA = dihedral(p1,p2,p3,p4) - phi_A0;
            uphi_B = (1-cos(dB)); dB = dihedral(p2,p3,p4,p5) - phi_B0;
            uphi_C = (1-cos(dC)); dC = dihedral(p3,p4,p5,p6) - phi_C0;
            pi = %f;
            """ % np.pi
        return energy_function

    def _numerical_torsion_integrand(self, phi, phi0, spring_constant, kt):
        """
        Integrand for the torsion (phi) restraints which will be integrated numerically for standard state correction

        Uses a periodic restraint where the magnitude force constant is matched to the standard harmonic by tailor
        series expansion around phi0.

        Domain is on [0, 2*pi]

        Parameters
        ----------
        phi : float or np.ndarray of float
            Torsion angle which will be integrated, units of radians
        phi0 : float
            Equilibrium torsion angle at which force of restraint often is 0, units of radians
        spring_constant : float
            Spring constant for this torsion in units of kJ/mol/nm**2
        kt : float
            Boltzmann Temperature of the thermodynamic state restraining the atoms = kB * T
            in units of kJ/mol

        Returns
        -------
        integrand : float
            Value of the integrated
        """
        def hav(angle):
            """Simple Haversine function for space later"""
            return (1-np.cos(angle))/2

        # The 1/2 from the K/2 is just folded into the hav() function
        return np.exp(-(spring_constant / kt) * hav(phi - phi0))


class RMSD(OpenMM73, ReceptorLigandRestraint):
    """Impose RMSD restraint on protein-ligand system.

    This restrains both protein and ligand using a flat-bottom RMSD restraint of the form:

    ``E = lambda_restraints * step(RMSD-RMSD0) * (K/2)*(RMSD-RMSD0)^2``

    Parameters
    ----------
    restrained_receptor_atoms : iterable of int, str, or None; Optional
        The indices of the receptor atoms to restrain, an MDTraj DSL expression, or a
        :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Select String <yank.Topography.select>`.
        Any number of receptor atoms can be selected.
        This can temporarily be left undefined, but ``determine_missing_parameters()``
        must be called before using the Restraint object. The same if a DSL
        expression or Topography region is provided.
        If no selection is given, all receptor atoms will be restrained.
        If an empty list is provided, no receptor atoms will be restrained.
        (default is None).
    restrained_ligand_atoms : iterable of int, str, or None; Optional
        The indices of the ligand atoms to restrain, an MDTraj DSL expression, or a
        :class:`Topography <yank.Topography>` region name,
        or :func:`Topography Select String <yank.Topography.select>`.
        Any number of ligand atoms can be selected
        This can temporarily be left undefined, but ``determine_missing_parameters()``
        must be called before using the Restraint object. The same if a DSL
        expression or Topography region is provided
        If no selection is given, all ligand atoms will be restrained.
        If an empty list is provided, no receptor atoms will be restrained.
        (default is None).
    K_RMSD : simtk.unit.Quantity, optional, default=0.6*kilocalories_per_mole/angstrom**2
        The spring constant (units compatible with kilocalories_per_mole/angstrom**2).
    RMSD0 : simtk.unit.Quantity, optional, default=2.0*angstrom
        The RMSD at which the restraint becomes nonzero.
    reference_sampler_state : openmmtools.states.SamplerState or None, Optional
        Reference sampler state with coordinates to use as the structure to align the RMSD to.
        The sampler state must have the same number of particles as the thermodynamic state you
        will apply this force to. This can temporarily be left undefined, but
        ``determine_missing_parameters()`` must be called before using the Restraint object.


    Attributes
    ----------
    restrained_receptor_atoms : list of int
        The indices of the restrained receptor atoms
    restrained_ligand_atoms : list of int
        The indices of the restrained_ligand_atoms

    Examples
    --------
    Create the ThermodynamicState.

    >>> from openmmtools import testsystems, states
    >>> system_container = testsystems.LysozymeImplicit()
    >>> system, positions = system_container.system, system_container.positions
    >>> thermodynamic_state = states.ThermodynamicState(system, 298*unit.kelvin)
    >>> sampler_state = states.SamplerState(positions)

    Identify ligand atoms. Topography automatically identify receptor atoms too.

    >>> from yank.yank import Topography
    >>> topography = Topography(system_container.topology, ligand_atoms=range(2603, 2621))

    Create a restraint

    >>> restraint = RMSD(restrained_receptor_atoms=[1335, 1339, 1397],
    ...                  restrained_ligand_atoms=[2609, 2607, 2606],
    ...                  K_RMSD=1.0*unit.kilocalories_per_mole/unit.angstrom**2,
    ...                  RMSD0=1*unit.angstroms)

    Find missing parameters

    >>> restraint.determine_missing_parameters(thermodynamic_state, sampler_state, topography)

    Get standard state correction.

    >>> correction = restraint.get_standard_state_correction(thermodynamic_state)

    """
    def __init__(self, restrained_receptor_atoms=None, restrained_ligand_atoms=None,
                 K_RMSD=None, RMSD0=None,
                 reference_sampler_state=None):
        super().__init__()
        self.restrained_receptor_atoms = restrained_receptor_atoms
        self.restrained_ligand_atoms = restrained_ligand_atoms
        self.K_RMSD = K_RMSD
        self.RMSD0 = RMSD0
        self.reference_sampler_state = reference_sampler_state

    # -------------------------------------------------------------------------
    # Public properties.
    # -------------------------------------------------------------------------

    class _RMSDRestrainedAtomsProperty(_RestrainedAtomsProperty):
        """
        Descriptor of restrained atoms.

        Extends `_RestrainedAtomsProperty` to handle single integers and strings.

        Extension allows individual atom lists to be defined as empty through the ``allowed_empty`` boolean
        """

        _MUST_COMPUTE_STRING = ('You are specifying {} {} atoms, '
                                'the final atoms will be chosen at from this set but you MUST '
                                'run "determine_missing_parameters"')

        def __init__(self, atoms_type, allowed_empty=False):
            self._allowed_empty = allowed_empty
            super().__init__(atoms_type)

        @methoddispatch
        def _validate_atoms(self, restrained_atoms):
            restrained_atoms = super()._validate_atoms(restrained_atoms)
            # TODO: Determine the minimum number of atoms needed for this restraint (can it be 0?)
            if len(restrained_atoms) < 3 and not (len(restrained_atoms) == 0 and self._allowed_empty):
                raise ValueError('At least three {} atoms are required to impose an '
                                 'RMSD restraint.'.format(self._atoms_type))
            return restrained_atoms

        @_validate_atoms.register(str)
        def _cast_atom_string(self, restrained_atoms):
            logger.warning(self._MUST_COMPUTE_STRING.format("a string for", self._atoms_type))
            return restrained_atoms

    restrained_receptor_atoms = _RMSDRestrainedAtomsProperty('receptor', allowed_empty=True)
    restrained_ligand_atoms = _RMSDRestrainedAtomsProperty('ligand')

    # -------------------------------------------------------------------------
    # Public methods.
    # -------------------------------------------------------------------------

    def restrain_state(self, thermodynamic_state):
        """Add the restraint force to the state's ``System``.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state holding the system to modify.

        """

        # Check if all parameters are defined.
        self._check_parameters_defined()

        # Merge receptor and ligand atoms in a single array for easy manipulation.
        restrained_atoms = self.restrained_receptor_atoms + self.restrained_ligand_atoms

        # Create RMSDForce CV for all restrained atoms
        rmsd_cv = openmm.RMSDForce(self.reference_sampler_state.positions, restrained_atoms)

        # Create an CustomCVForce
        energy_expression = 'lambda_restraints * step(dRMSD) * (K_RMSD/2)*dRMSD^2; dRMSD = (RMSD-RMSD0);'
        energy_expression += 'K_RMSD = %f;' % self.K_RMSD.value_in_unit_system(unit.md_unit_system)
        energy_expression += 'RMSD0 = %f;' % self.RMSD0.value_in_unit_system(unit.md_unit_system)
        restraint_force = openmm.CustomCVForce(energy_expression)
        restraint_force.addCollectiveVariable('RMSD', rmsd_cv)
        restraint_force.addGlobalParameter('lambda_restraints', 1.0)

        # Get a copy of the system of the ThermodynamicState, modify it and set it back.
        system = thermodynamic_state.system
        self._add_force_in_separate_group(system, restraint_force)
        thermodynamic_state.system = system

    def get_standard_state_correction(self, thermodynamic_state):
        """Return the standard state correction.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state.

        Returns
        -------
        DeltaG : float
           Computed standard-state correction in dimensionless units (kT).

        """
        # TODO: Compute standard state correction
        return 0.0

    def determine_missing_parameters(self, thermodynamic_state, sampler_state, topography):
        """Set reference positions for RMSD restraint.

        Future iterations of this feature will introduce the ability to extract
        equilibrium parameters and spring constants from a short simulation.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state.
        sampler_state : openmmtools.states.SamplerState, optional
            The sampler state holding the positions of all atoms to be used as reference
        topography : yank.Topography, optional
            The topography with labeled receptor and ligand atoms.

        """

        # Pick out the restrained atoms
        if not self._are_restrained_atoms_defined:
            self._pick_restrained_atoms(topography)

        # We determine automatically only the parameters that have been left undefined.
        def _assign_if_undefined(attr_name, attr_value):
            """Assign value to self.name only if it is None."""
            if getattr(self, attr_name) is None:
                setattr(self, attr_name, attr_value)

        # Set spring constants uniformly, as in Ref [1] Table 1 caption.
        _assign_if_undefined('K_RMSD', 0.6 * unit.kilocalories_per_mole / unit.angstrom**2)
        _assign_if_undefined('RMSD0', 2.0 * unit.angstroms)

        # Write restraint parameters
        msg = 'restraint parameters:\n'
        for parameter_name, parameter_value in self._parameters.items():
            msg += '{0:24s} : {1}\n'.format(parameter_name, parameter_value)
        logger.debug(msg)

        # Assign the sampler state
        _assign_if_undefined('reference_sampler_state', sampler_state)

    # -------------------------------------------------------------------------
    # Internal-usage
    # -------------------------------------------------------------------------

    def _check_parameters_defined(self):
        """Raise an exception there are still parameters undefined."""
        if not self._are_restrained_atoms_defined:
            raise RestraintParameterError('Undefined restrained atoms. Please run `determine_missing_parameters`')
        if None in [self.K_RMSD, self.RMSD0]:
            raise RestraintParameterError("Undefined RMSD parameters. Please run `determine_missing_parameters`")
        if self.reference_sampler_state is None:
            raise RestraintParameterError("No reference configuration/structure as defined, "
                                          "Please run `determine_missing_parameters`")

    @property
    def _are_restrained_atoms_defined(self):
        """Check if the restrained atoms are defined well enough to make a restraint"""
        for atoms in [self.restrained_receptor_atoms, self.restrained_ligand_atoms]:
            # Atoms should be a list or None at this point due to the _RestrainedAtomsProperty class
            if not self._are_single_atoms_defined(atoms):
                return False
        return True

    @staticmethod
    def _are_single_atoms_defined(atom_list):
        """Check that a set of atoms matches the specific format we expect"""
        if isinstance(atom_list, list):
            return True
        return False

    def _pick_restrained_atoms(self, topography):
        """Select the restrained atoms to use for this system"""
        atom_selector = _AtomSelector(topography)
        for atom_word, top_key in zip(["restrained_ligand_atoms", "restrained_receptor_atoms"],
                                      ["ligand_atoms",            "receptor_atoms"]):
            atoms = getattr(self, atom_word)
            if self._are_single_atoms_defined(atoms):
                continue
            defined_atoms = atom_selector.compute_atom_intersect(atoms, top_key)
            setattr(self, atom_word, defined_atoms)


if __name__ == '__main__':
    import doctest
    doctest.testmod()
    # doctest.run_docstring_examples(Harmonic, globals())
