# Source code for propertyestimator.protocols.simulation

    """
    A collection of protocols for running molecular simulations.
    """
    import logging
    import math
    import os
    import shutil
    import threading
    import traceback
    from enum import Enum

    import numpy as np
    import yaml
    from simtk import openmm, unit as simtk_unit
    from simtk.openmm import app

    from propertyestimator import unit
    from propertyestimator.thermodynamics import ThermodynamicState, Ensemble
    from propertyestimator.utils.exceptions import PropertyEstimatorException
    from propertyestimator.utils.openmm import setup_platform_with_resources, openmm_quantity_to_pint, \
        pint_quantity_to_openmm, disable_pbc
    from propertyestimator.utils.quantities import EstimatedQuantity
    from propertyestimator.utils.statistics import StatisticsArray, ObservableType
    from propertyestimator.utils.utils import temporarily_change_directory, safe_unlink
    from propertyestimator.workflow.decorators import protocol_input, protocol_output, MergeBehaviour
    from propertyestimator.workflow.plugins import register_calculation_protocol
    from propertyestimator.workflow.protocols import BaseProtocol

    @register_calculation_protocol()
    class RunEnergyMinimisation(BaseProtocol):
        """A protocol to minimise the potential energy of a system.
        """

        @protocol_input(str)
        def input_coordinate_file(self):
            """The coordinates to minimise."""
            pass

        @protocol_input(unit.Quantity)
        def tolerance(self):
            """The energy tolerance to which the system should be minimized."""
            pass

        @protocol_input(int)
        def max_iterations(self):
            """The maximum number of iterations to perform.  If this is 0,
            minimization is continued until the results converge without regard
            to how many iterations it takes."""
            pass

        @protocol_input(str)
        def system_path(self):
            """The path to the XML system object which defines the forces present in the system."""
            pass

        @protocol_input(bool)
        def enable_pbc(self):
            """If true, periodic boundary conditions will be enabled."""
            pass

        @protocol_output(str)
        def output_coordinate_file(self):
            """The file path to the minimised coordinates."""
            pass

        def __init__(self, protocol_id):

            super().__init__(protocol_id)

            self._input_coordinate_file = None

            self._system_path = None
            self._system = None

            self._enable_pbc = True

            self._tolerance = 10*unit.kilojoules / unit.mole
            self._max_iterations = 0

            self._output_coordinate_file = None

        def execute(self, directory, available_resources):

            logging.info('Minimising energy: ' + self.id)

            platform = setup_platform_with_resources(available_resources)

            input_pdb_file = app.PDBFile(self._input_coordinate_file)

            with open(self._system_path, 'rb') as file:
                self._system = openmm.XmlSerializer.deserialize(file.read().decode())

            if not self._enable_pbc:

                for force_index in range(self._system.getNumForces()):

                    force = self._system.getForce(force_index)

                    if not isinstance(force, openmm.NonbondedForce):
                        continue

                    force.setNonbondedMethod(0)  # NoCutoff = 0, NonbondedMethod.CutoffNonPeriodic = 1

            # TODO: Expose the constraint tolerance
            integrator = openmm.VerletIntegrator(0.002 * simtk_unit.picoseconds)
            simulation = app.Simulation(input_pdb_file.topology, self._system, integrator, platform)

            box_vectors = input_pdb_file.topology.getPeriodicBoxVectors()

            if box_vectors is None:
                box_vectors = simulation.system.getDefaultPeriodicBoxVectors()

            simulation.context.setPeriodicBoxVectors(*box_vectors)
            simulation.context.setPositions(input_pdb_file.positions)

            simulation.minimizeEnergy(pint_quantity_to_openmm(self._tolerance), self._max_iterations)

            positions = simulation.context.getState(getPositions=True).getPositions()

            self._output_coordinate_file = os.path.join(directory, 'minimised.pdb')

            with open(self._output_coordinate_file, 'w+') as minimised_file:
                app.PDBFile.writeFile(simulation.topology, positions, minimised_file)

            logging.info('Energy minimised: ' + self.id)

            return self._get_output_dictionary()

    @register_calculation_protocol()
    class RunOpenMMSimulation(BaseProtocol):
        """Performs a molecular dynamics simulation in a given ensemble using
        an OpenMM backend.
        """

        @protocol_input(int, merge_behavior=MergeBehaviour.GreatestValue)
        def steps(self):
            """The number of timesteps to evolve the system by."""
            pass

        @protocol_input(unit.Quantity, merge_behavior=MergeBehaviour.SmallestValue)
        def thermostat_friction(self):
            """The thermostat friction coefficient."""
            pass

        @protocol_input(unit.Quantity, merge_behavior=MergeBehaviour.SmallestValue)
        def timestep(self):
            """The timestep to evolve the system by at each step."""
            pass

        @protocol_input(int, merge_behavior=MergeBehaviour.SmallestValue)
        def output_frequency(self):
            """The frequency with which to write to the output statistics and trajectory files."""
            pass

        @protocol_input(Ensemble)
        def ensemble(self):
            """The thermodynamic ensemble to simulate in."""
            pass

        @protocol_input(ThermodynamicState)
        def thermodynamic_state(self):
            """The thermodynamic conditions to simulate under"""
            pass

        @protocol_input(str)
        def input_coordinate_file(self):
            """The file path to the starting coordinates."""
            pass

        @protocol_input(str)
        def system_path(self):
            """A path to the XML system object which defines the forces present in the system."""
            pass

        @protocol_input(bool)
        def enable_pbc(self):
            """If true, periodic boundary conditions will be enabled."""
            pass

        @protocol_input(bool)
        def save_rolling_statistics(self):
            """If True, the statisitics file will be written to every
            `output_frequency` number of steps, rather than just once
            at the end of the simulation.

            Notes
            -----
            In future when either saving the statistics to file has been
            optimised, or an option for the frequency to save to the file
            has been added, this option will be removed.
            """
            pass

        @protocol_input(bool)
        def allow_gpu_platforms(self):
            """If true, OpenMM will be allowed to run using
            a GPU if available, otherwise it will be constrained
            to only using CPUs."""
            pass

        @protocol_input(bool)
        def high_precision(self):
            """If true, OpenMM will be run using a platform with
            high precision settings. This will be the Reference
            platform when only a CPU is available, or double
            precision mode when a GPU is available."""
            pass

        @protocol_output(str)
        def output_coordinate_file(self):
            """The file path to the coordinates of the final system configuration."""
            pass

        @protocol_output(str)
        def trajectory_file_path(self):
            """The file path to the trajectory sampled during the simulation."""
            pass

        @protocol_output(str)
        def statistics_file_path(self):
            """The file path to the statistics sampled during the simulation."""
            pass

        def __init__(self, protocol_id):

            super().__init__(protocol_id)

            self._steps = 1000

            self._thermostat_friction = 1.0 / unit.picoseconds
            self._timestep = 0.002 * unit.picoseconds

            self._output_frequency = 500000

            self._ensemble = Ensemble.NPT

            self._input_coordinate_file = None
            self._thermodynamic_state = None

            self._system_path = None

            self._enable_pbc = True

            self._save_rolling_statistics = True

            self._allow_gpu_platforms = True
            self._high_precision = False

            self._output_coordinate_file = None

            self._trajectory_file_path = None
            self._statistics_file_path = None

            # Keep a track of the file names used for temporary working files.
            self._temporary_statistics_path = None
            self._temporary_trajectory_path = None

            self._checkpoint_path = None

            self._context = None
            self._integrator = None

        def execute(self, directory, available_resources):

            # We handle most things in OMM units here.
            temperature = pint_quantity_to_openmm(self._thermodynamic_state.temperature)
            pressure = pint_quantity_to_openmm(None if self._ensemble == Ensemble.NVT else
                                               self._thermodynamic_state.pressure)

            if temperature is None:

                return PropertyEstimatorException(directory=directory,
                                                  message='A temperature must be set to perform '
                                                          'a simulation in any ensemble')

            if Ensemble(self._ensemble) == Ensemble.NPT and pressure is None:

                return PropertyEstimatorException(directory=directory,
                                                  message='A pressure must be set to perform an NPT simulation')

            if Ensemble(self._ensemble) == Ensemble.NPT and self._enable_pbc is False:

                return PropertyEstimatorException(directory=directory,
                                                  message='PBC must be enabled when running in the NPT ensemble.')

            logging.info('Performing a simulation in the ' + str(self._ensemble) + ' ensemble: ' + self.id)

            # Clean up any temporary files from previous (possibly failed)
            # simulations.
            self._temporary_statistics_path = os.path.join(directory, 'temp_statistics.csv')
            self._temporary_trajectory_path = os.path.join(directory, 'temp_trajectory.dcd')

            safe_unlink(self._temporary_statistics_path)
            safe_unlink(self._temporary_trajectory_path)

            self._checkpoint_path = os.path.join(directory, 'checkpoint.xml')

            # Set up the output file paths
            self._trajectory_file_path = os.path.join(directory, 'trajectory.dcd')
            self._statistics_file_path = os.path.join(directory, 'statistics.csv')

            # Set up the simulation objects.
            if self._context is None or self._integrator is None:

                self._context, self._integrator = self._setup_simulation_objects(temperature,
                                                                                 pressure,
                                                                                 available_resources)

            result = self._simulate(directory, temperature, pressure, self._context, self._integrator)
            return result

        def _setup_simulation_objects(self, temperature, pressure, available_resources):
            """Initializes the objects needed to perform the simulation.
            This comprises of a context, and an integrator.

            Parameters
            ----------
            temperature: simtk.unit.Quantity
                The temperature to run the simulation at.
            pressure: simtk.unit.Quantity
                The pressure to run the simulation at.
            available_resources: ComputeResources
                The resources available to run on.

            Returns
            -------
            simtk.openmm.Context
                The created openmm context which takes advantage
                of the available compute resources.
            openmmtools.integrators.LangevinIntegrator
                The Langevin integrator which will propogate
                the simulation.
            """

            import openmmtools
            from simtk.openmm import XmlSerializer

            # Create a platform with the correct resources.
            if not self._allow_gpu_platforms:

                from propertyestimator.backends import ComputeResources
                available_resources = ComputeResources(available_resources.number_of_threads)

            platform = setup_platform_with_resources(available_resources, self._high_precision)

            # Load in the system object from the provided xml file.
            with open(self._system_path, 'r') as file:
                system = XmlSerializer.deserialize(file.read())

            # Disable the periodic boundary conditions if requested.
            if not self._enable_pbc:

                disable_pbc(system)
                pressure = None

            # Use the openmmtools ThermodynamicState object to help
            # set up a system which contains the correct barostat if
            # one should be present.
            openmm_state = openmmtools.states.ThermodynamicState(system=system,
                                                                 temperature=temperature,
                                                                 pressure=pressure)

            system = openmm_state.get_system(remove_thermostat=True)

            # Set up the integrator.
            thermostat_friction = pint_quantity_to_openmm(self._thermostat_friction)
            timestep = pint_quantity_to_openmm(self._timestep)

            integrator = openmmtools.integrators.LangevinIntegrator(temperature=temperature,
                                                                    collision_rate=thermostat_friction,
                                                                    timestep=timestep)

            # Create the simulation context.
            context = openmm.Context(system, integrator, platform)

            # Initialize the context with the correct positions etc.
            if os.path.isfile(self._checkpoint_path):

                # Load the simulation state from a checkpoint file.
                logging.info(f'Loading the checkpoint from {self._checkpoint_path}.')

                with open(self._checkpoint_path, 'r') as file:
                    checkpoint_state = XmlSerializer.deserialize(file.read())

                context.setState(checkpoint_state)

            else:

                logging.info(f'No checkpoint file was found at {self._checkpoint_path}.')

                # Populate the simulation object from the starting input files.
                input_pdb_file = app.PDBFile(self._input_coordinate_file)

                if self._enable_pbc:

                    # Optionally set up the box vectors.
                    box_vectors = input_pdb_file.topology.getPeriodicBoxVectors()

                    if box_vectors is None:

                        raise ValueError('The input file must contain box vectors '
                                         'when running with PBC.')

                    context.setPeriodicBoxVectors(*box_vectors)

                context.setPositions(input_pdb_file.positions)
                context.setVelocitiesToTemperature(temperature)

            return context, integrator

        def _write_statistics_array(self, raw_statistics, current_step, temperature, pressure,
                                    degrees_of_freedom, total_mass):
            """Appends a set of statistics to an existing `statistics_array`.
            Those statistics are potential energy, kinetic energy, total energy,
            volume, density and reduced potential.

            Parameters
            ----------
            raw_statistics: dict of ObservableType and numpy.ndarray
                A dictionary of potential energies (kJ/mol), kinetic
                energies (kJ/mol) and volumes (angstrom**3).
            current_step: int
                The index of the current step.
            temperature: simtk.unit.Quantity
                The temperature the system is being simulated at.
            pressure: simtk.unit.Quantity
                The pressure the system is being simulated at.
            degrees_of_freedom: int
                The number of degrees of freedom the system has.
            total_mass: simtk.unit.Quantity
                The total mass of the system.

            Returns
            -------
            StatisticsArray
                The statistics array with statistics appended.
            """
            temperature = openmm_quantity_to_pint(temperature)
            pressure = openmm_quantity_to_pint(pressure)

            beta = 1.0 / (unit.boltzmann_constant * temperature)

            raw_potential_energies = raw_statistics[ObservableType.PotentialEnergy][0:current_step + 1]
            raw_kinetic_energies = raw_statistics[ObservableType.KineticEnergy][0:current_step + 1]
            raw_volumes = raw_statistics[ObservableType.Volume][0:current_step + 1]

            potential_energies = raw_potential_energies * unit.kilojoules / unit.mole
            kinetic_energies = raw_kinetic_energies * unit.kilojoules / unit.mole
            volumes = raw_volumes * unit.angstrom ** 3

            # Calculate the instantaneous temperature, taking account the
            # systems degrees of freedom.
            temperatures = 2.0 * kinetic_energies / (degrees_of_freedom * unit.molar_gas_constant)

            # Calculate the systems enthalpy and reduced potential.
            total_energies = potential_energies + kinetic_energies
            enthalpies = None

            reduced_potentials = potential_energies / unit.avogadro_number

            if pressure is not None:

                pv_terms = pressure * volumes

                reduced_potentials += pv_terms
                enthalpies = total_energies + pv_terms * unit.avogadro_number

            reduced_potentials = (beta * reduced_potentials) * unit.dimensionless

            # Calculate the systems density.
            densities = total_mass / (volumes * unit.avogadro_number)

            statistics_array = StatisticsArray()

            statistics_array[ObservableType.PotentialEnergy] = potential_energies
            statistics_array[ObservableType.KineticEnergy] = kinetic_energies
            statistics_array[ObservableType.TotalEnergy] = total_energies
            statistics_array[ObservableType.Temperature] = temperatures
            statistics_array[ObservableType.Volume] = volumes
            statistics_array[ObservableType.Density] = densities
            statistics_array[ObservableType.ReducedPotential] = reduced_potentials

            if enthalpies is not None:
                statistics_array[ObservableType.Enthalpy] = enthalpies

            statistics_array.to_pandas_csv(self._temporary_statistics_path)

        def _simulate(self, directory, temperature, pressure, context, integrator):
            """Performs the simulation using a given context
            and integrator.

            Parameters
            ----------
            directory: str
                The directory the trajectory is being run in.
            temperature: simtk.unit.Quantity
                The temperature to run the simulation at.
            pressure: simtk.unit.Quantity
                The pressure to run the simulation at.
            context: simtk.openmm.Context
                The OpenMM context to run with.
            integrator: simtk.openmm.Integrator
                The integrator to evolve the simulation with.
            """

            # Build the reporters which we will use to report the state
            # of the simulation.
            input_pdb_file = app.PDBFile(self._input_coordinate_file)
            topology = input_pdb_file.topology

            with open(os.path.join(directory, 'input.pdb'), 'w+') as configuration_file:
                app.PDBFile.writeFile(input_pdb_file.topology, input_pdb_file.positions, configuration_file)

            # Make a copy of the existing trajectory to append to if one already exists.
            append_trajectory = False

            if os.path.isfile(self._trajectory_file_path):

                shutil.copyfile(self._trajectory_file_path, self._temporary_trajectory_path)
                append_trajectory = True

            elif os.path.isfile(self._temporary_trajectory_path):
                os.unlink(self._temporary_trajectory_path)

            if append_trajectory:
                trajectory_file_object = open(self._temporary_trajectory_path, 'r+b')
            else:
                trajectory_file_object = open(self._temporary_trajectory_path, 'w+b')

            trajectory_dcd_object = app.DCDFile(trajectory_file_object,
                                                topology,
                                                integrator.getStepSize(),
                                                0,
                                                self._output_frequency,
                                                append_trajectory)

            expected_number_of_statistics = math.ceil(self.steps / self.output_frequency)

            raw_statistics = {
                ObservableType.PotentialEnergy: np.zeros(expected_number_of_statistics),
                ObservableType.KineticEnergy: np.zeros(expected_number_of_statistics),
                ObservableType.Volume: np.zeros(expected_number_of_statistics),
            }

            # Define any constants needed for extracting system statistics
            # Compute the instantaneous temperature of degrees of freedom.
            # This snipped is taken from the build in OpenMM `StateDataReporter`
            system = context.getSystem()

            degrees_of_freedom = sum([3 for i in range(system.getNumParticles()) if
                                      system.getParticleMass(i) > 0 * simtk_unit.dalton])

            degrees_of_freedom -= system.getNumConstraints()

            if any(type(system.getForce(i)) == openmm.CMMotionRemover for i in range(system.getNumForces())):
                degrees_of_freedom -= 3

            total_mass = 0.0 * simtk_unit.dalton

            for i in range(system.getNumParticles()):
                total_mass += system.getParticleMass(i)

            total_mass = openmm_quantity_to_pint(total_mass)

            # Perform the simulation.
            current_step_count = 0
            current_step = 0

            result = None

            try:

                while current_step_count < self.steps:

                    steps_to_take = min(self._output_frequency, self.steps - current_step_count)
                    integrator.step(steps_to_take)

                    state = context.getState(getPositions=True,
                                             getEnergy=True,
                                             getVelocities=False,
                                             getForces=False,
                                             getParameters=False,
                                             enforcePeriodicBox=self.enable_pbc)

                    # Write out the current frame of the trajectory.
                    trajectory_dcd_object.writeModel(positions=state.getPositions(),
                                                     periodicBoxVectors=state.getPeriodicBoxVectors())

                    # Write out the energies and system statistics.
                    raw_statistics[ObservableType.PotentialEnergy][current_step] = \
                        state.getPotentialEnergy().value_in_unit(simtk_unit.kilojoules_per_mole)
                    raw_statistics[ObservableType.KineticEnergy][current_step] = \
                        state.getKineticEnergy().value_in_unit(simtk_unit.kilojoules_per_mole)
                    raw_statistics[ObservableType.Volume][current_step] = \
                        state.getPeriodicBoxVolume().value_in_unit(simtk_unit.angstrom ** 3)

                    if self._save_rolling_statistics:

                        self._write_statistics_array(raw_statistics, current_step, temperature,
                                                     pressure, degrees_of_freedom, total_mass)

                    current_step_count += steps_to_take
                    current_step += 1

            except Exception as e:

                formatted_exception = f'{traceback.format_exception(None, e, e.__traceback__)}'

                result = PropertyEstimatorException(directory=directory,
                                                    message=f'The simulation failed unexpectedly: '
                                                            f'{formatted_exception}')

            # Create a checkpoint file.
            state = context.getState(getPositions=True,
                                     getEnergy=True,
                                     getVelocities=True,
                                     getForces=True,
                                     getParameters=True,
                                     enforcePeriodicBox=self.enable_pbc)

            state_xml = openmm.XmlSerializer.serialize(state)

            with open(self._checkpoint_path, 'w') as file:
                file.write(state_xml)

            # Make sure to close the open trajectory stream.
            trajectory_file_object.close()

            # Save the final statistics
            self._write_statistics_array(raw_statistics, current_step, temperature,
                                         pressure, degrees_of_freedom, total_mass)

            if isinstance(result, PropertyEstimatorException):
                return result

            # Move the trajectory and statistics files to their
            # final location.
            os.replace(self._temporary_trajectory_path, self._trajectory_file_path)

            if not os.path.isfile(self._statistics_file_path):
                os.replace(self._temporary_statistics_path, self._statistics_file_path)
            else:

                existing_statistics = StatisticsArray.from_pandas_csv(self._statistics_file_path)
                current_statistics = StatisticsArray.from_pandas_csv(self._temporary_statistics_path)

                concatenated_statistics = StatisticsArray.join(existing_statistics,
                                                               current_statistics)

                concatenated_statistics.to_pandas_csv(self._statistics_file_path)

            # Save out the final positions.
            final_state = context.getState(getPositions=True)
            positions = final_state.getPositions()
            topology.setPeriodicBoxVectors(final_state.getPeriodicBoxVectors())

            self._output_coordinate_file = os.path.join(directory, 'output.pdb')

            with open(self._output_coordinate_file, 'w+') as configuration_file:
                app.PDBFile.writeFile(topology, positions, configuration_file)

            logging.info(f'Simulation performed in the {str(self._ensemble)} ensemble: {self._id}')
            return self._get_output_dictionary()

    @register_calculation_protocol()
    class BaseYankProtocol(BaseProtocol):
        """An abstract base class for protocols which will performs a set of alchemical
        free energy simulations using the YANK framework.

        Protocols which inherit from this base must implement the abstract `_get_yank_options`
        methods.
        """

        @protocol_input(ThermodynamicState)
        def thermodynamic_state(self):
            """The state at which to run the calculations."""
            pass

        @protocol_input(int, merge_behavior=MergeBehaviour.GreatestValue)
        def number_of_iterations(self):
            """The number of YANK iterations to perform."""
            pass

        @protocol_input(int, merge_behavior=MergeBehaviour.GreatestValue)
        def steps_per_iteration(self):
            """The number of steps per YANK iteration to perform."""
            pass

        @protocol_input(int, merge_behavior=MergeBehaviour.SmallestValue)
        def checkpoint_interval(self):
            """The number of iterations between saving YANK checkpoint files."""
            pass

        @protocol_input(unit.Quantity, merge_behavior=MergeBehaviour.SmallestValue)
        def timestep(self):
            """The length of the timestep to take."""
            pass

        @protocol_input(str)
        def force_field_path(self):
            """The path to the force field to use for the calculations"""
            pass

        @protocol_input(bool)
        def verbose(self):
            """Controls whether or not to run YANK at high verbosity."""
            pass

        @protocol_output(EstimatedQuantity)
        def estimated_free_energy(self):
            """The estimated free energy value and its uncertainty returned
            by YANK."""
            pass

        def __init__(self, protocol_id):
            """Constructs a new BaseYankProtocol object."""

            super().__init__(protocol_id)

            self._thermodynamic_state = None
            self._timestep = 2 * unit.femtosecond

            self._number_of_iterations = 1

            self._steps_per_iteration = 500
            self._checkpoint_interval = 50

            self._force_field_path = None

            self._verbose = False

            self._estimated_free_energy = None

        def _get_options_dictionary(self, available_resources):
            """Returns a dictionary of options which will be serialized
            to a yaml file and passed to YANK.

            Parameters
            ----------
            available_resources: ComputeResources
                The resources available to execute on.

            Returns
            -------
            dict of str and Any
                A yaml compatible dictionary of YANK options.
            """

            from openforcefield.utils import quantity_to_string

            platform_name = 'CPU'

            if available_resources.number_of_gpus > 0:

                # A platform which runs on GPUs has been requested.
                from propertyestimator.backends import ComputeResources
                toolkit_enum = ComputeResources.GPUToolkit(available_resources.preferred_gpu_toolkit)

                # A platform which runs on GPUs has been requested.
                platform_name = 'CUDA' if toolkit_enum == ComputeResources.GPUToolkit.CUDA else \
                                                          ComputeResources.GPUToolkit.OpenCL

            return {
                'verbose': self._verbose,
                'output_dir': '.',

                'temperature': quantity_to_string(pint_quantity_to_openmm(self._thermodynamic_state.temperature)),
                'pressure': quantity_to_string(pint_quantity_to_openmm(self._thermodynamic_state.pressure)),

                'minimize': True,

                'default_number_of_iterations': self._number_of_iterations,
                'default_nsteps_per_iteration': self._steps_per_iteration,
                'checkpoint_interval': self._checkpoint_interval,

                'default_timestep': quantity_to_string(pint_quantity_to_openmm(self._timestep)),

                'annihilate_electrostatics': True,
                'annihilate_sterics': False,

                'platform': platform_name
            }

        def _get_solvent_dictionary(self):
            """Returns a dictionary of the solvent which will be serialized
            to a yaml file and passed to YANK. In most cases, this should
            just be passing force field settings over, such as PME settings.

            Returns
            -------
            dict of str and Any
                A yaml compatible dictionary of YANK solvents.
            """
            from openforcefield.typing.engines.smirnoff.forcefield import ForceField
            force_field = ForceField(self._force_field_path)

            charge_method = force_field.get_parameter_handler('Electrostatics').method

            if charge_method.lower() != 'pme':
                raise ValueError('Currently only PME electrostatics are supported.')

            return {'default': {
                'nonbonded_method': charge_method,
            }}

        def _get_system_dictionary(self):
            """Returns a dictionary of the system which will be serialized
            to a yaml file and passed to YANK. Only a single system may be
            specified.

            Returns
            -------
            dict of str and Any
                A yaml compatible dictionary of YANK systems.
            """
            raise NotImplementedError()

        def _get_protocol_dictionary(self):
            """Returns a dictionary of the protocol which will be serialized
            to a yaml file and passed to YANK. Only a single protocol may be
            specified.

            Returns
            -------
            dict of str and Any
                A yaml compatible dictionary of a YANK protocol.
            """
            raise NotImplementedError()

        def _get_experiments_dictionary(self):
            """Returns a dictionary of the experiments which will be serialized
            to a yaml file and passed to YANK. Only a single experiment may be
            specified.

            Returns
            -------
            dict of str and Any
                A yaml compatible dictionary of a YANK experiment.
            """

            system_dictionary = self._get_system_dictionary()
            system_key = next(iter(system_dictionary))

            protocol_dictionary = self._get_protocol_dictionary()
            protocol_key = next(iter(protocol_dictionary))

            return {
                'system': system_key,
                'protocol': protocol_key
            }

        def _get_full_input_dictionary(self, available_resources):
            """Returns a dictionary of the full YANK inputs which will be serialized
            to a yaml file and passed to YANK

            Parameters
            ----------
            available_resources: ComputeResources
                The resources available to execute on.

            Returns
            -------
            dict of str and Any
                A yaml compatible dictionary of a YANK input file.
            """

            return {
                'options': self._get_options_dictionary(available_resources),

                'solvents': self._get_solvent_dictionary(),

                'systems': self._get_system_dictionary(),
                'protocols': self._get_protocol_dictionary(),

                'experiments': self._get_experiments_dictionary()
            }

        @staticmethod
        def _extract_trajectory(checkpoint_path, output_trajectory_path):
            """Extracts the stored trajectory of the 'initial' state from a
            yank `.nc` checkpoint file and stores it to disk as a `.dcd` file.

            Parameters
            ----------
            checkpoint_path: str
                The path to the yank `.nc` file
            output_trajectory_path: str
                The path to store the extracted trajectory at.
            """

            from yank.analyze import extract_trajectory

            mdtraj_trajectory = extract_trajectory(checkpoint_path, state_index=0, image_molecules=True)
            mdtraj_trajectory.save_dcd(output_trajectory_path)

        @staticmethod
        def _run_yank(directory, available_resources):
            """Runs YANK within the specified directory which contains a `yank.yaml`
            input file.

            Parameters
            ----------
            directory: str
                The directory within which to run yank.

            Returns
            -------
            simtk.unit.Quantity
                The free energy returned by yank.
            simtk.unit.Quantity
                The uncertainty in the free energy returned by yank.
            """

            from yank.experiment import ExperimentBuilder
            from yank.analyze import ExperimentAnalyzer

            with temporarily_change_directory(directory):

                # Set the default properties on the desired platform
                # before calling into yank.
                setup_platform_with_resources(available_resources)

                exp_builder = ExperimentBuilder('yank.yaml')
                exp_builder.run_experiments()

                analyzer = ExperimentAnalyzer('experiments')
                output = analyzer.auto_analyze()

                free_energy = output['free_energy']['free_energy_diff_unit']
                free_energy_uncertainty = output['free_energy']['free_energy_diff_error_unit']

            return free_energy, free_energy_uncertainty

        @staticmethod
        def _run_yank_as_process(queue, directory, available_resources):
            """A wrapper around the `_run_yank` method which takes
            a `multiprocessing.Queue` as input, thereby allowing it
            to be launched from a separate process and still return
            it's output back to the main process.

            Parameters
            ----------
            queue: multiprocessing.Queue
                The queue object which will communicate with the
                launched process.
            directory: str
                The directory within which to run yank.

            Returns
            -------
            simtk.unit.Quantity
                The free energy returned by yank.
            simtk.unit.Quantity
                The uncertainty in the free energy returned by yank.
            str, optional
                The stringified errors which occurred on the other process,
                or `None` if no exceptions were raised.
            """

            free_energy = None
            free_energy_uncertainty = None

            error = None

            try:
                free_energy, free_energy_uncertainty = BaseYankProtocol._run_yank(directory, available_resources)
            except Exception as e:
                error = traceback.format_exception(None, e, e.__traceback__)

            queue.put((free_energy, free_energy_uncertainty, error))

        def execute(self, directory, available_resources):

            yaml_filename = os.path.join(directory, 'yank.yaml')

            # Create the yank yaml input file from a dictionary of options.
            with open(yaml_filename, 'w') as file:
                yaml.dump(self._get_full_input_dictionary(available_resources), file)

            # Yank is not safe to be called from anything other than the main thread.
            # If the current thread is not detected as the main one, then yank should
            # be spun up in a new process which should itself be safe to run yank in.
            if threading.current_thread() is threading.main_thread():
                logging.info('Launching YANK in the main thread.')
                free_energy, free_energy_uncertainty = self._run_yank(directory, available_resources)
            else:

                from multiprocessing import Process, Queue

                logging.info('Launching YANK in a new process.')

                # Create a queue to pass the results back to the main process.
                queue = Queue()
                # Create the process within which yank will run.
                process = Process(target=BaseYankProtocol._run_yank_as_process, args=[queue, directory,
                                                                                      available_resources])

                # Start the process and gather back the output.
                process.start()
                free_energy, free_energy_uncertainty, error = queue.get()
                process.join()

                if error is not None:
                    return PropertyEstimatorException(directory, error)

            self._estimated_free_energy = EstimatedQuantity(openmm_quantity_to_pint(free_energy),
                                                            openmm_quantity_to_pint(free_energy_uncertainty),
                                                            self._id)

            return self._get_output_dictionary()

    @register_calculation_protocol()
    class LigandReceptorYankProtocol(BaseYankProtocol):
        """An abstract base class for protocols which will performs a set of
        alchemical free energy simulations using the YANK framework.

        Protocols which inherit from this base must implement the abstract
        `_get_*_dictionary` methods.
        """

        class RestraintType(Enum):
            """The types of ligand restraints available within yank.
            """
            Harmonic = 'Harmonic'
            FlatBottom = 'FlatBottom'

        @protocol_input(str)
        def ligand_residue_name(self):
            """The residue name of the ligand."""
            pass

        @protocol_input(str)
        def receptor_residue_name(self):
            """The residue name of the receptor."""
            pass

        @protocol_input(str)
        def solvated_ligand_coordinates(self):
            """The file path to the solvated ligand coordinates."""
            pass

        @protocol_input(str)
        def solvated_ligand_system(self):
            """The file path to the solvated ligand system object."""
            pass

        @protocol_input(str)
        def solvated_complex_coordinates(self):
            """The file path to the solvated complex coordinates."""
            pass

        @protocol_input(str)
        def solvated_complex_system(self):
            """The file path to the solvated complex system object."""
            pass

        @protocol_input(bool)
        def apply_restraints(self):
            """Determines whether the ligand should be explicitly restrained to
            the receptor in order to stop the ligand from temporarily unbinding.
            """
            pass

        @protocol_input(RestraintType)
        def restraint_type(self):
            """The type of ligand restraint applied, provided that `apply_restraints`
            is `True`"""
            pass

        @protocol_output(str)
        def solvated_ligand_trajectory_path(self):
            """The file path to the generated ligand trajectory."""
            pass

        @protocol_output(str)
        def solvated_complex_trajectory_path(self):
            """The file path to the generated ligand trajectory."""
            pass

        def __init__(self, protocol_id):
            """Constructs a new LigandReceptorYankProtocol object."""

            super().__init__(protocol_id)

            self._ligand_residue_name = None
            self._receptor_residue_name = None

            self._solvated_ligand_coordinates = None
            self._solvated_ligand_system = None

            self._solvated_complex_coordinates = None
            self._solvated_complex_system = None

            self._local_ligand_coordinates = 'ligand.pdb'
            self._local_ligand_system = 'ligand.xml'

            self._local_complex_coordinates = 'complex.pdb'
            self._local_complex_system = 'complex.xml'

            self._solvated_ligand_trajectory_path = None
            self._solvated_complex_trajectory_path = None

            self._apply_restraints = True
            self._restraint_type = LigandReceptorYankProtocol.RestraintType.Harmonic

        def _get_system_dictionary(self):

            solvent_dictionary = self._get_solvent_dictionary()
            solvent_key = next(iter(solvent_dictionary))

            host_guest_dictionary = {
                'phase1_path': [self._local_complex_system, self._local_complex_coordinates],
                'phase2_path': [self._local_ligand_system, self._local_ligand_coordinates],

                'ligand_dsl': f'resname {self._ligand_residue_name}',
                'solvent': solvent_key
            }

            return {'host-guest': host_guest_dictionary}

        def _get_protocol_dictionary(self):

            absolute_binding_dictionary = {
                'complex': {'alchemical_path': 'auto'},
                'solvent': {'alchemical_path': 'auto'}
            }

            return {'absolute_binding_dictionary': absolute_binding_dictionary}

        def _get_experiments_dictionary(self):

            experiments_dictionary = super(LigandReceptorYankProtocol, self)._get_experiments_dictionary()

            if self._apply_restraints:

                experiments_dictionary['restraint'] = {
                    'restrained_ligand_atoms': f'(resname {self._ligand_residue_name}) and (mass > 1.5)',
                    'restrained_receptor_atoms': f'(resname {self._receptor_residue_name}) and (mass > 1.5)',

                    'type': self._restraint_type.value
                }

            return experiments_dictionary

        def execute(self, directory, available_resources):

            # Because of quirks in where Yank looks files while doing temporary
            # directory changes, we need to copy the coordinate files locally so
            # they are correctly found.
            shutil.copyfile(self._solvated_ligand_coordinates, os.path.join(directory, self._local_ligand_coordinates))
            shutil.copyfile(self._solvated_ligand_system, os.path.join(directory, self._local_ligand_system))

            shutil.copyfile(self._solvated_complex_coordinates, os.path.join(directory, self._local_complex_coordinates))
            shutil.copyfile(self._solvated_complex_system, os.path.join(directory, self._local_complex_system))

            result = super(LigandReceptorYankProtocol, self).execute(directory, available_resources)

            if isinstance(result, PropertyEstimatorException):
                return result

            ligand_yank_path = os.path.join(directory, 'experiments', 'solvent.nc')
            complex_yank_path = os.path.join(directory, 'experiments', 'complex.nc')

            self._solvated_ligand_trajectory_path = os.path.join(directory, 'ligand.dcd')
            self._solvated_complex_trajectory_path = os.path.join(directory, 'complex.dcd')

            self._extract_trajectory(ligand_yank_path, self._solvated_ligand_trajectory_path)
            self._extract_trajectory(complex_yank_path, self._solvated_complex_trajectory_path)

            return self._get_output_dictionary()