"""
A collection of protocols for performing free energy calculations
using the YANK package.
"""
import logging
import os
import shutil
import threading
import traceback
from enum import Enum

import yaml

from propertyestimator import unit
from propertyestimator.forcefield import SmirnoffForceFieldSource
from propertyestimator.substances import Substance
from propertyestimator.thermodynamics import ThermodynamicState
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.utils import temporarily_change_directory
from propertyestimator.workflow.decorators import protocol_input, protocol_output, InequalityMergeBehaviour, UNDEFINED
from propertyestimator.workflow.plugins import register_calculation_protocol
from propertyestimator.workflow.protocols import BaseProtocol


@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.
    """

    thermodynamic_state = protocol_input(
        docstring='The state at which to run the calculations.',
        type_hint=ThermodynamicState,
        default_value=UNDEFINED
    )

    number_of_equilibration_iterations = protocol_input(
        docstring='The number of iterations used for equilibration before production run. '
                  'Only post-equilibration iterations are written to file.',
        type_hint=int, merge_behavior=InequalityMergeBehaviour.LargestValue,
        default_value=1
    )
    number_of_iterations = protocol_input(
        docstring='The number of YANK iterations to perform.',
        type_hint=int, merge_behavior=InequalityMergeBehaviour.LargestValue,
        default_value=5000
    )
    steps_per_iteration = protocol_input(
        docstring='The number of steps per YANK iteration to perform.',
        type_hint=int, merge_behavior=InequalityMergeBehaviour.LargestValue,
        default_value=500
    )
    checkpoint_interval = protocol_input(
        docstring='The number of iterations between saving YANK checkpoint files.',
        type_hint=int, merge_behavior=InequalityMergeBehaviour.SmallestValue,
        default_value=50
    )

    timestep = protocol_input(
        docstring='The length of the timestep to take.',
        type_hint=unit.Quantity, merge_behavior=InequalityMergeBehaviour.SmallestValue,
        default_value=2 * unit.femtosecond
    )

    verbose = protocol_input(
        docstring='Controls whether or not to run YANK at high verbosity.',
        type_hint=bool,
        default_value=False
    )
    setup_only = protocol_input(
        docstring='If true, YANK will only create and validate the setup files, '
                  'but not actually run any simulations. This argument is mainly '
                  'only to be used for testing purposes.',
        type_hint=bool,
        default_value=False
    )

    estimated_free_energy = protocol_output(
        docstring='The estimated free energy value and its uncertainty '
                  'returned by YANK.',
        type_hint=EstimatedQuantity
    )

    @staticmethod
    def _get_residue_names_from_role(substances, coordinate_path, role):
        """Returns a list of all of the residue names of
        components which have been assigned a given role.

        Parameters
        ----------
        substances: list of Substance
            The substances which contains the components.
        coordinate_path: str
            The path to the coordinates which describe the systems
            topology.
        role: Substance.ComponentRole
            The role of the component to identify.

        Returns
        -------
        set of str
            The identified residue names.
        """

        from simtk.openmm import app
        from openforcefield.topology import Molecule, Topology

        if role == Substance.ComponentRole.Undefined:
            return 'all'

        unique_molecules = [Molecule.from_smiles(component.smiles) for
                            substance in substances for
                            component in substance.components]

        openmm_topology = app.PDBFile(coordinate_path).topology
        topology = Topology.from_openmm(openmm_topology, unique_molecules)

        # Determine the smiles of all molecules in the system. We need to use
        # the toolkit to re-generate the smiles as later we will compare these
        # against more toolkit generated smiles.
        components = [component for substance in substances for
                      component in substance.components if
                      component.role == role]

        component_smiles = [Molecule.from_smiles(component.smiles).to_smiles() for
                            component in components]

        residue_names = set()

        all_openmm_atoms = list(openmm_topology.atoms())

        # Find the resiude names of the molecules which have the correct
        # role.
        for topology_molecule in topology.topology_molecules:

            molecule_smiles = topology_molecule.reference_molecule.to_smiles()

            if molecule_smiles not in component_smiles:
                continue

            molecule_residue_names = set([all_openmm_atoms[topology_atom.topology_atom_index].residue.name for
                                          topology_atom in topology_molecule.atoms])

            assert len(molecule_residue_names) == 1
            residue_names.update(molecule_residue_names)

        return residue_names

    @staticmethod
    def _get_dsl_from_role(substances, coordinate_path, role):
        """Returns an MDTraj DSL string which identifies those
        atoms which belong to components flagged with a specific
        role.

        Parameters
        ----------
        substances: list of Substance
            The substances which contains the components.
        coordinate_path: str
            The path to the coordinates which describe the systems
            topology.
        role: Substance.ComponentRole
            The role of the component to identify.

        Returns
        -------
        str
            The DSL string.
        """

        residue_names = BaseYankProtocol._get_residue_names_from_role(substances, coordinate_path, role)

        dsl_string = ' or '.join([f'resname {residue_name}' for residue_name in residue_names])
        return dsl_string

    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,

            'number_of_equilibration_iterations': self.number_of_equilibration_iterations,
            '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_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),

            '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, setup_only):
        """Runs YANK within the specified directory which contains a `yank.yaml`
        input file.

        Parameters
        ----------
        directory: str
            The directory within which to run yank.
        available_resources: ComputeResources
            The compute resources available to yank.
        setup_only: bool
            If true, YANK will only create and validate the setup files,
            but not actually run any simulations. This argument is mainly
            only to be used for testing purposes.

        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

        from simtk import unit as simtk_unit

        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')

            if setup_only is True:
                return 0.0 * simtk_unit.kilojoule_per_mole, 0.0 * simtk_unit.kilojoule_per_mole

            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, setup_only):
        """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.
        available_resources: ComputeResources
            The compute resources available to yank.
        setup_only: bool
            If true, YANK will only create and validate the setup files,
            but not actually run any simulations. This argument is mainly
            only to be used for testing purposes.

        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,
                                                                              setup_only)
        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, sort_keys=False)

        setup_only = self.setup_only

        # 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, setup_only)
        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, setup_only])

            # 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):
    """A protocol for performing ligand-receptor alchemical free energy
    calculations using the YANK framework.
    """

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

    ligand_residue_name = protocol_input(
        docstring='The residue name of the ligand.',
        type_hint=str,
        default_value=UNDEFINED
    )
    receptor_residue_name = protocol_input(
        docstring='The residue name of the receptor.',
        type_hint=str,
        default_value=UNDEFINED
    )

    solvated_ligand_coordinates = protocol_input(
        docstring='The file path to the solvated ligand coordinates.',
        type_hint=str,
        default_value=UNDEFINED
    )
    solvated_ligand_system = protocol_input(
        docstring='The file path to the solvated ligand system object.',
        type_hint=str,
        default_value=UNDEFINED
    )

    solvated_complex_coordinates = protocol_input(
        docstring='The file path to the solvated complex coordinates.',
        type_hint=str,
        default_value=UNDEFINED
    )
    solvated_complex_system = protocol_input(
        docstring='The file path to the solvated complex system object.',
        type_hint=str,
        default_value=UNDEFINED
    )

    force_field_path = protocol_input(
        docstring='The path to the force field which defines the charge method '
                  'to use for the calculation.',
        type_hint=str,
        default_value=UNDEFINED
    )

    apply_restraints = protocol_input(
        docstring='Determines whether the ligand should be explicitly restrained to the '
                  'receptor in order to stop the ligand from temporarily unbinding.',
        type_hint=bool,
        default_value=True
    )
    restraint_type = protocol_input(
        docstring='The type of ligand restraint applied, provided that `apply_restraints` '
                  'is `True`',
        type_hint=RestraintType,
        default_value=RestraintType.Harmonic
    )

    solvated_ligand_trajectory_path = protocol_output(
        docstring='The file path to the generated ligand trajectory.',
        type_hint=str
    )
    solvated_complex_trajectory_path = protocol_output(
        docstring='The file path to the generated ligand trajectory.',
        type_hint=str
    )

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

        super().__init__(protocol_id)

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

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

    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.
        """

        with open(self.force_field_path, 'r') as file:
            force_field_source = SmirnoffForceFieldSource.parse_json(file.read())

        force_field = force_field_source.to_force_field()
        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):

        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 _get_full_input_dictionary(self, available_resources):

        full_dictionary = super(LigandReceptorYankProtocol, self)._get_full_input_dictionary(available_resources)
        full_dictionary['solvents'] = self._get_solvent_dictionary()

        return full_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

        if self.setup_only:
            return self._get_output_dictionary()

        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()


@register_calculation_protocol()
class SolvationYankProtocol(BaseYankProtocol):
    """A protocol for performing solvation alchemical free energy
    calculations using the YANK framework.

    This protocol can be used for box solvation free energies (setting
    the `solvent_1` input to the solvent of interest and setting
    `solvent_2` as an empty `Substance`) or transfer free energies (setting
    both the `solvent_1` and `solvent_2` inputs to different solvents).
    """

    solute = protocol_input(
        docstring='The substance describing the composition of '
                  'the solute. This should include the solute '
                  'molecule as well as any counter ions.',
        type_hint=Substance,
        default_value=UNDEFINED
    )

    solvent_1 = protocol_input(
        docstring='The substance describing the composition of '
                  'the first solvent.',
        type_hint=Substance,
        default_value=UNDEFINED
    )
    solvent_2 = protocol_input(
        docstring='The substance describing the composition of '
                  'the second solvent.',
        type_hint=Substance,
        default_value=UNDEFINED
    )

    solvent_1_coordinates = protocol_input(
        docstring='The file path to the coordinates of the solute embedded in the '
                  'first solvent.',
        type_hint=str,
        default_value=UNDEFINED
    )
    solvent_1_system = protocol_input(
        docstring='The file path to the system object of the solute embedded in the '
                  'first solvent.',
        type_hint=str,
        default_value=UNDEFINED
    )

    solvent_2_coordinates = protocol_input(
        docstring='The file path to the coordinates of the solute embedded in the '
                  'second solvent.',
        type_hint=str,
        default_value=UNDEFINED
    )
    solvent_2_system = protocol_input(
        docstring='The file path to the system object of the solute embedded in the '
                  'second solvent.',
        type_hint=str,
        default_value=UNDEFINED
    )

    electrostatic_lambdas_1 = protocol_input(
        docstring='The list of electrostatic alchemical states that YANK should sample at. '
                  'These values will be passed to the YANK `lambda_electrostatics` option. '
                  'If no option is set, YANK will use `trailblaze` algorithm to determine '
                  'this option automatically.',
        type_hint=list,
        optional=True,
        default_value=UNDEFINED
    )
    steric_lambdas_1 = protocol_input(
        docstring='The list of steric alchemical states that YANK should sample at. '
                  'These values will be passed to the YANK `lambda_sterics` option. '
                  'If no option is set, YANK will use `trailblaze` algorithm to determine '
                  'this option automatically.',
        type_hint=list,
        optional=True,
        default_value=UNDEFINED
    )
    electrostatic_lambdas_2 = protocol_input(
        docstring='The list of electrostatic alchemical states that YANK should sample at. '
                  'These values will be passed to the YANK `lambda_electrostatics` option. '
                  'If no option is set, YANK will use `trailblaze` algorithm to determine '
                  'this option automatically.',
        type_hint=list,
        optional=True,
        default_value=UNDEFINED
    )
    steric_lambdas_2 = protocol_input(
        docstring='The list of steric alchemical states that YANK should sample at. '
                  'These values will be passed to the YANK `lambda_sterics` option. '
                  'If no option is set, YANK will use `trailblaze` algorithm to determine '
                  'this option automatically.',
        type_hint=list,
        optional=True,
        default_value=UNDEFINED
    )

    solvent_1_trajectory_path = protocol_output(
        docstring='The file path to the trajectory of the solute in the '
                  'first solvent.',
        type_hint=str
    )
    solvent_2_trajectory_path = protocol_output(
        docstring='The file path to the trajectory of the solute in the '
                  'second solvent.',
        type_hint=str
    )

    def __init__(self, protocol_id):
        super().__init__(protocol_id)

        self._local_solvent_1_coordinates = 'solvent_1.pdb'
        self._local_solvent_1_system = 'solvent_1.xml'

        self._local_solvent_2_coordinates = 'solvent_2.pdb'
        self._local_solvent_2_system = 'solvent_2.xml'

    def _get_system_dictionary(self):

        solvent_1_dsl = self._get_dsl_from_role([self.solute, self.solvent_1],
                                                self.solvent_1_coordinates,
                                                Substance.ComponentRole.Solvent)

        solvent_2_dsl = self._get_dsl_from_role([self.solute, self.solvent_2],
                                                self.solvent_2_coordinates,
                                                Substance.ComponentRole.Solvent)

        full_solvent_dsl_components = []

        if len(solvent_1_dsl) > 0:
            full_solvent_dsl_components.append(solvent_1_dsl)
        if len(solvent_2_dsl) > 0:
            full_solvent_dsl_components.append(solvent_2_dsl)

        solvation_system_dictionary = {
            'phase1_path': [self._local_solvent_1_system, self._local_solvent_1_coordinates],
            'phase2_path': [self._local_solvent_2_system, self._local_solvent_2_coordinates],

            'solvent_dsl': ' or '.join(full_solvent_dsl_components)
        }

        return {'solvation-system': solvation_system_dictionary}

    def _get_protocol_dictionary(self):

        solvent_1_protocol_dictionary = {
            'lambda_electrostatics': self.electrostatic_lambdas_1,
            'lambda_sterics': self.steric_lambdas_1
        }

        if self.electrostatic_lambdas_1 == UNDEFINED and self.steric_lambdas_1 == UNDEFINED:

            solvent_1_protocol_dictionary = 'auto'

        elif ((self.electrostatic_lambdas_1 != UNDEFINED and self.steric_lambdas_1 == UNDEFINED) or
              (self.electrostatic_lambdas_1 == UNDEFINED and self.steric_lambdas_1 != UNDEFINED)):

            raise ValueError('Either both of `electrostatic_lambdas_1` and '
                             '`steric_lambdas_1` must be set, or neither '
                             'must be set.')

        solvent_2_protocol_dictionary = {
            'lambda_electrostatics': self.electrostatic_lambdas_2,
            'lambda_sterics': self.steric_lambdas_2
        }
        
        if self.electrostatic_lambdas_2 == UNDEFINED and self.steric_lambdas_2 == UNDEFINED:

            solvent_2_protocol_dictionary = 'auto'

        elif ((self.electrostatic_lambdas_2 != UNDEFINED and self.steric_lambdas_2 == UNDEFINED) or
              (self.electrostatic_lambdas_2 == UNDEFINED and self.steric_lambdas_2 != UNDEFINED)):

            raise ValueError('Either both of `electrostatic_lambdas_2` and '
                             '`steric_lambdas_2` must be set, or neither '
                             'must be set.')

        protocol_dictionary = {
            'solvent1': {'alchemical_path': solvent_1_protocol_dictionary},
            'solvent2': {'alchemical_path': solvent_2_protocol_dictionary}
        }

        return {'solvation-protocol': protocol_dictionary}

    def execute(self, directory, available_resources):

        from simtk.openmm import XmlSerializer

        solute_components = [component for component in self.solute.components if
                             component.role == Substance.ComponentRole.Solute]

        solvent_1_components = [component for component in self.solvent_1.components if
                                component.role == Substance.ComponentRole.Solvent]

        solvent_2_components = [component for component in self.solvent_2.components if
                                component.role == Substance.ComponentRole.Solvent]

        if len(solute_components) != 1:
            return PropertyEstimatorException(directory, 'There must only be a single component marked as a solute.')
        if len(solvent_1_components) == 0 and len(solvent_2_components) == 0:
            return PropertyEstimatorException(directory, 'At least one of the solvents must not be vacuum.')

        # 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.solvent_1_coordinates, os.path.join(directory, self._local_solvent_1_coordinates))
        shutil.copyfile(self.solvent_1_system, os.path.join(directory, self._local_solvent_1_system))

        shutil.copyfile(self.solvent_2_coordinates, os.path.join(directory, self._local_solvent_2_coordinates))
        shutil.copyfile(self.solvent_2_system, os.path.join(directory, self._local_solvent_2_system))

        # Disable the pbc of the any solvents which should be treated
        # as vacuum.
        vacuum_system_path = None

        if len(solvent_1_components) == 0:
            vacuum_system_path = self._local_solvent_1_system
        elif len(solvent_2_components) == 0:
            vacuum_system_path = self._local_solvent_2_system

        if vacuum_system_path is not None:

            logging.info(f'Disabling the periodic boundary conditions in {vacuum_system_path} '
                         f'by setting the cutoff type to NoCutoff')

            with open(os.path.join(directory, vacuum_system_path), 'r') as file:
                vacuum_system = XmlSerializer.deserialize(file.read())

            disable_pbc(vacuum_system)

            with open(os.path.join(directory, vacuum_system_path), 'w') as file:
                file.write(XmlSerializer.serialize(vacuum_system))

        # Set up the yank input file.
        result = super(SolvationYankProtocol, self).execute(directory, available_resources)

        if isinstance(result, PropertyEstimatorException):
            return result

        if self.setup_only:
            return self._get_output_dictionary()

        solvent_1_yank_path = os.path.join(directory, 'experiments', 'solvent1.nc')
        solvent_2_yank_path = os.path.join(directory, 'experiments', 'solvent2.nc')

        self.solvent_1_trajectory_path = os.path.join(directory, 'solvent1.dcd')
        self.solvent_2_trajectory_path = os.path.join(directory, 'solvent2.dcd')

        self._extract_trajectory(solvent_1_yank_path, self.solvent_1_trajectory_path)
        self._extract_trajectory(solvent_2_yank_path, self.solvent_2_trajectory_path)

        return self._get_output_dictionary()
