#!/usr/bin/env python

# =============================================================================================
# MODULE DOCSTRING
# =============================================================================================

"""
Pipeline
========

Utility functions to help setting up Yank configurations.

"""

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

import collections
import copy
import inspect
import itertools
import json
import logging
import os
import re
import sys

import mdtraj
import mpiplus
import numpy as np
import openmmtools as mmtools
import openmoltools as moltools
import yaml
from pdbfixer import PDBFixer
from simtk import openmm, unit
from simtk.openmm.app import PDBFile


from . import utils

logger = logging.getLogger(__name__)


# ==============================================================================
# Utility functions
# ==============================================================================

def compute_squared_distances(molecule1_positions, molecule2_positions):
    """Compute the squared distances between the atoms of two molecules.

    All the positions must be expressed in the same unit of measure.

    Parameters
    ----------
    molecule1_positions : numpy.ndarray
        An Nx3 array where, N is the number of atoms, containing the positions of
        the atoms of molecule1.
    molecule2_positions : numpy.ndarray
        An Mx3 array where, M is the number of atoms, containing the positions of
        the atoms of the molecule2.

    Returns
    -------
    squared_distances : numpy.ndarray
        An NxM array of squared distances. distances_squared[i][j] is the squared
        distance between atom i of molecule1 and atom j of molecule 2.

    """
    squared_distances = np.array([((molecule2_positions - atom_pos)**2).sum(1)
                                  for atom_pos in molecule1_positions])
    return squared_distances


def compute_min_dist(mol_positions, *args):
    """Compute the minimum distance between a molecule and a set of other molecules.

    All the positions must be expressed in the same unit of measure.

    Parameters
    ----------
    mol_positions : numpy.ndarray
        An Nx3 array where, N is the number of atoms, containing the positions of
        the atoms of the molecule for which we want to compute the minimum distance
        from the others

    Other parameters
    ----------------
    args
        A series of numpy.ndarrays containing the positions of the atoms of the other
        molecules

    Returns
    -------
    min_dist : float
        The minimum distance between ``mol_positions`` and the other set of positions

    """
    for argmol_positions in args:
        # Compute squared distances.  Each row is an array of distances
        # from a mol_positions atom to all argmol_positions atoms.
        distances2 = compute_squared_distances(mol_positions, argmol_positions)

        # Find closest atoms and their distance
        min_idx = np.unravel_index(distances2.argmin(), distances2.shape)
        try:
            min_dist = min(min_dist, np.sqrt(distances2[min_idx]))
        except UnboundLocalError:
            min_dist = np.sqrt(distances2[min_idx])
    return min_dist


def compute_min_max_dist(mol_positions, *args):
    """Compute minimum and maximum distances between a molecule and a set of
    other molecules.

    All the positions must be expressed in the same unit of measure.

    Parameters
    ----------
    mol_positions : numpy.ndarray
        An Nx3 array where, N is the number of atoms, containing the positions of
        the atoms of the molecule for which we want to compute the minimum distance
        from the others

    Other parameters
    ----------------
    args
        A series of numpy.ndarrays containing the positions of the atoms of the other
        molecules

    Returns
    -------
    min_dist : float
        The minimum distance between mol_positions and the atoms of the other positions
    max_dist : float
        The maximum distance between mol_positions and the atoms of the other positions

    Examples
    --------
    >>> mol1_pos = np.array([[-1, -1, -1], [1, 1, 1]], np.float)
    >>> mol2_pos = np.array([[2, 2, 2], [2, 4, 5]], np.float)  # determine min dist
    >>> mol3_pos = np.array([[3, 3, 3], [3, 4, 5]], np.float)  # determine max dist
    >>> min_dist, max_dist = compute_min_max_dist(mol1_pos, mol2_pos, mol3_pos)
    >>> min_dist == np.linalg.norm(mol1_pos[1] - mol2_pos[0])
    True
    >>> max_dist == np.linalg.norm(mol1_pos[1] - mol3_pos[1])
    True

    """
    min_dist = None

    for argmol_positions in args:
        # Compute squared distances of all atoms. Each row is an array of distances
        # from an atom in argmol_positions to all the atoms in mol_positions.
        distances2 = compute_squared_distances(argmol_positions, mol_positions)

        # Find distances of each arg_pos atom to mol_positions
        distances2 = np.amin(distances2, axis=1)

        # Find closest and distant atom
        if min_dist is None:
            min_dist = np.sqrt(distances2.min())
            max_dist = np.sqrt(distances2.max())
        else:
            min_dist = min(min_dist, np.sqrt(distances2.min()))
            max_dist = max(max_dist, np.sqrt(distances2.max()))

    return min_dist, max_dist


def compute_radius_of_gyration(positions):
    """
    Compute the radius of gyration of the specified coordinate set.

    Parameters
    ----------
    positions : simtk.unit.Quantity with units compatible with angstrom
       The coordinate set (natoms x 3) for which the radius of gyration is to be computed.

    Returns
    -------
    radius_of_gyration : simtk.unit.Quantity with units compatible with angstrom
       The radius of gyration

    """
    unit = positions.unit

    # Get dimensionless receptor positions.
    x = positions / unit

    # Get dimensionless restrained atom coordinate.
    xref = x.mean(0)
    xref = np.reshape(xref, (1,3)) # (1,3) array

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

    # Compute std dev of distances from restrained atom.
    radius_of_gyration = distances.std() * unit

    return radius_of_gyration


def compute_net_charge(system, atom_indices):
    """Compute the total net charge of a subset of atoms in the system.

    Parameters
    ----------
    system : simtk.openmm.System
        The system object containing the atoms of interest.
    atom_indices : list of int
        Indices of the atoms of interest.

    Returns
    -------
    net_charge : int
        Total net charge as the sum of the partial charges of the atoms.

    """
    atom_indices = set(atom_indices)  # convert to set to speed up searching
    net_charge = 0.0 * unit.elementary_charge
    for force_index in range(system.getNumForces()):
        force = system.getForce(force_index)
        if isinstance(force, openmm.NonbondedForce):
            for particle_index in range(force.getNumParticles()):
                if particle_index in atom_indices:
                    net_charge += force.getParticleParameters(particle_index)[0]
                    atom_indices.remove(particle_index)
    assert len(atom_indices) == 0
    net_charge = int(round(net_charge / unit.elementary_charge))
    return net_charge


def find_alchemical_counterions(system, topography, region_name):
    """Return the atom indices of the ligand or solute counter ions.

    In periodic systems, the solvation box needs to be neutral, and
    if the decoupled molecule is charged, it will cause trouble. This
    can be used to find a set of ions in the system that neutralize
    the molecule, so that the solvation box will remain neutral all
    the time.

    Parameters
    ----------
    system : simtk.openmm.System
        The system object containing the atoms of interest.
    topography : yank.Topography
        The topography object holding the indices of the ions and the
        ligand (for binding free energy) or solute (for transfer free
        energy).
    region_name : str
        The region name in the topography (e.g. "ligand_atoms") for
        which to find counter ions.

    Returns
    -------
    counterions_indices : list of int
        The list of atom indices in the system of the counter ions
        neutralizing the region.

    Raises
    ------
    ValueError
        If the topography region has no atoms, or if it impossible
        to neutralize the region with the ions in the system.

    """
    # Check whether we need to find counterions of ligand or solute.
    atom_indices = getattr(topography, region_name)
    if len(atom_indices) == 0:
        raise ValueError("Cannot find counterions for region {}. "
                         "The region has no atoms.")

    # If the net charge of alchemical atoms is 0, we don't need counterions.
    mol_net_charge = compute_net_charge(system, atom_indices)
    logger.debug('{} net charge: {}'.format(region_name, mol_net_charge))
    if mol_net_charge == 0:
        return []

    # Find net charge of all ions in the system.
    ions_net_charges = [(ion_id, compute_net_charge(system, [ion_id]))
                        for ion_id in topography.ions_atoms]
    topology = topography.topology
    ions_names_charges = [(topology.atom(ion_id).residue.name, ion_net_charge)
                          for ion_id, ion_net_charge in ions_net_charges]
    logger.debug('Ions net charges: {}'.format(ions_names_charges))

    # Find minimal subset of counterions whose charges sums to -mol_net_charge.
    for n_ions in range(1, len(ions_net_charges) + 1):
        for ion_subset in itertools.combinations(ions_net_charges, n_ions):
            counterions_indices, counterions_charges = zip(*ion_subset)
            if sum(counterions_charges) == -mol_net_charge:
                return counterions_indices

    # We couldn't find any subset of counterions neutralizing the region.
    raise ValueError('Impossible to find a solution for region {}. '
                     'Net charge: {}, system ions: {}.'.format(
        region_name, mol_net_charge, ions_names_charges))


# See Amber manual Table 4.1 http://ambermd.org/doc12/Amber15.pdf
_OPENMM_TO_TLEAP_PBRADII = {'HCT': 'mbondi', 'OBC1': 'mbondi2', 'OBC2': 'mbondi2',
                            'GBn': 'bondi', 'GBn2': 'mbondi3'}


def get_leap_recommended_pbradii(implicit_solvent):
    """Return the recommended PBradii setting for LeAP.

    Parameters
    ----------
    implicit_solvent : str
        The implicit solvent model.

    Returns
    -------
    pbradii : str or object
        The LeAP recommended PBradii for the model.

    Raises
    ------
    ValueError
        If the implicit solvent model is not supported by OpenMM.

    Examples
    --------
    >>> get_leap_recommended_pbradii('OBC2')
    'mbondi2'
    >>> from simtk.openmm.app import HCT
    >>> get_leap_recommended_pbradii(HCT)
    'mbondi'

    """
    try:
        return _OPENMM_TO_TLEAP_PBRADII[str(implicit_solvent)]
    except KeyError:
        raise ValueError('Implicit solvent {} is not supported.'.format(implicit_solvent))


_NONPERIODIC_NONBONDED_METHODS = [openmm.app.NoCutoff, openmm.app.CutoffNonPeriodic]


def create_system(parameters_file, box_vectors, create_system_args, system_options):
    """Create and return an OpenMM system.

    Parameters
    ----------
    parameters_file : simtk.openmm.app.AmberPrmtopFile or GromacsTopFile
        The file used to create they system.
    box_vectors : list of Vec3
        The default box vectors of the system will be set to this value.
    create_system_args : dict of str
        The kwargs accepted by the ``createSystem()`` function of the ``parameters_file``.
    system_options : dict
        The kwargs to forward to ``createSystem()``.

    Returns
    -------
    system : simtk.openmm.System
        The system created.

    """
    # Prepare createSystem() options
    # OpenMM adopts camel case convention so we need to change the options format.
    # Then we filter system options according to specific createSystem() args
    system_options = {utils.underscore_to_camelcase(key): value
                      for key, value in system_options.items()}
    system_options = {arg: system_options[arg] for arg in create_system_args
                      if arg in system_options}

    # Determine if this will be an explicit or implicit solvent simulation.
    if box_vectors is not None:
        is_periodic = True
    else:
        is_periodic = False

    # Adjust nonbondedMethod
    # TODO: Ensure that selected method is appropriate.
    if 'nonbondedMethod' not in system_options:
        if is_periodic:
            system_options['nonbondedMethod'] = openmm.app.CutoffPeriodic
        else:
            system_options['nonbondedMethod'] = openmm.app.NoCutoff

    # Check for solvent configuration inconsistencies
    # TODO: Check to make sure both files agree on explicit/implicit.
    err_msg = ''
    nonbonded_method = system_options['nonbondedMethod']
    if is_periodic:
        if 'implicitSolvent' in system_options and system_options['implicitSolvent'] is not None:
            err_msg = 'Found periodic box in positions file and implicitSolvent specified.'
        if nonbonded_method in _NONPERIODIC_NONBONDED_METHODS:
            err_msg = ('Found periodic box in positions file but '
                       'nonbondedMethod is {}'.format(nonbonded_method))
    else:
        if nonbonded_method not in _NONPERIODIC_NONBONDED_METHODS:
            err_msg = ('nonbondedMethod {} is periodic but could not '
                       'find periodic box in positions file.'.format(nonbonded_method))
    if len(err_msg) != 0:
        logger.error(err_msg)
        raise RuntimeError(err_msg)

    # Create system and update box vectors (if needed)
    system = parameters_file.createSystem(removeCMMotion=False, **system_options)
    if is_periodic:
        system.setDefaultPeriodicBoxVectors(*box_vectors)

    return system


def read_system_files(positions_file_path, parameters_file_path, system_options,
                      gromacs_include_dir=None, charmm_parameter_files=None):
    """Create a Yank arguments for a phase from system files.

    Parameters
    ----------
    positions_file_path : str
        Path to system position file (e.g. 'complex.inpcrd/.gro/.pdb').
    parameters_file_path : str
        Path to system parameters file (e.g. 'complex.prmtop/.top/.xml/.psf').
    system_options : dict
        ``system_options[phase]`` is a a dictionary containing options to
        pass to ``createSystem()``. If the parameters file is an OpenMM
        system in XML format, this will be ignored.
    gromacs_include_dir : str, optional
        Path to directory in which to look for other files included
        from the gromacs top file.
    charmm_parameter_files : str, optional
        Path to additional parameter files

    Returns
    -------
    system : simtk.openmm.System
        The OpenMM System built from the given files.
    topology : openmm.app.Topology
        The OpenMM Topology built from the given files.
    sampler_state : openmmtools.states.SamplerState
        The sampler state containing the positions of the atoms.

    """
    # Load system files
    parameters_file_extension = os.path.splitext(parameters_file_path)[1]

    # Read OpenMM XML and PDB files.
    if parameters_file_extension == '.xml':
        logger.debug("xml: {}".format(parameters_file_path))
        logger.debug("pdb: {}".format(positions_file_path))

        positions_file = openmm.app.PDBFile(positions_file_path)
        parameters_file = positions_file  # Needed for topology.
        with open(parameters_file_path, 'r') as f:
            serialized_system = f.read()

        system = openmm.XmlSerializer.deserialize(serialized_system)
        box_vectors = positions_file.topology.getPeriodicBoxVectors()
        if box_vectors is None:
            box_vectors = system.getDefaultPeriodicBoxVectors()

    # Read Amber prmtop and inpcrd files.
    elif parameters_file_extension == '.prmtop':
        logger.debug("prmtop: {}".format(parameters_file_path))
        logger.debug("inpcrd: {}".format(positions_file_path))

        parameters_file = openmm.app.AmberPrmtopFile(parameters_file_path)
        positions_file = openmm.app.AmberInpcrdFile(positions_file_path)
        box_vectors = positions_file.boxVectors
        create_system_args = set(inspect.getargspec(openmm.app.AmberPrmtopFile.createSystem).args)

        system = create_system(parameters_file, box_vectors, create_system_args, system_options)

    # Read Gromacs top and gro files.
    elif parameters_file_extension == '.top':
        logger.debug("top: {}".format(parameters_file_path))
        logger.debug("gro: {}".format(positions_file_path))

        positions_file = openmm.app.GromacsGroFile(positions_file_path)

        # gro files must contain box vectors, so we must determine whether system
        # is non-periodic or not from provided nonbonded options
        # WARNING: This uses the private API for GromacsGroFile, and may break.
        if ('nonbonded_method' in system_options and
                system_options['nonbonded_method'] in _NONPERIODIC_NONBONDED_METHODS):
            logger.info('nonbonded_method = {}, so removing periodic box vectors '
                        'from gro file'.format(system_options['nonbonded_method']))
            for frame, box_vectors in enumerate(positions_file._periodicBoxVectors):
                positions_file._periodicBoxVectors[frame] = None

        box_vectors = positions_file.getPeriodicBoxVectors()
        parameters_file = openmm.app.GromacsTopFile(parameters_file_path,
                                                    periodicBoxVectors=box_vectors,
                                                    includeDir=gromacs_include_dir)
        create_system_args = set(inspect.getargspec(openmm.app.GromacsTopFile.createSystem).args)

        system = create_system(parameters_file, box_vectors, create_system_args, system_options)

    # Read CHARMM format psf and pdb files
    elif parameters_file_extension == '.psf':
        logger.debug("psf: {}".format(parameters_file_path))
        logger.debug("pdb: {}".format(positions_file_path))

        parameters_file = openmm.app.CharmmPsfFile(parameters_file_path)
        positions_file = openmm.app.PDBFile(positions_file_path)
        params = openmm.app.CharmmParameterSet(*charmm_parameter_files)

        box_vectors = positions_file.topology.getPeriodicBoxVectors()
        if box_vectors is None:
            box_vectors = system.getDefaultPeriodicBoxVectors()

        parameters_file.setBox(box_vectors[0][0], box_vectors[1][1], box_vectors[2][2])
        create_system_args = set(inspect.getargspec(openmm.app.CharmmPsfFile.createSystem).args)
        system_options['params'] = params
        system = create_system(parameters_file, box_vectors, create_system_args, system_options)
            
    # Unsupported file format.
    else:
        raise ValueError('Unsupported format for parameter file {}'.format(parameters_file_extension))

    # Store numpy positions and create SamplerState.
    positions = positions_file.getPositions(asNumpy=True)
    sampler_state = mmtools.states.SamplerState(positions=positions, box_vectors=box_vectors)

    # Check to make sure number of atoms match between prmtop and inpcrd.
    n_atoms_system = system.getNumParticles()
    n_atoms_positions = positions.shape[0]
    if n_atoms_system != n_atoms_positions:
        err_msg = "Atom number mismatch: {} has {} atoms; {} has {} atoms.".format(
            parameters_file_path, n_atoms_system, positions_file_path, n_atoms_positions)
        logger.error(err_msg)
        raise RuntimeError(err_msg)

    return system, parameters_file.topology, sampler_state


# =============================================================================
# SETUP PIPELINE UTILITY FUNCTIONS
# =============================================================================

# Map the OpenMM-style name for a solvent to the tleap
# name compatible with the solvateBox command.
_OPENMM_LEAP_SOLVENT_MODELS_MAP = {
    'tip3p': 'TIP3PBOX',
    'tip3pfb': 'TIP3PFBOX',
    'tip4pew': 'TIP4PEWBOX',
    'tip5p': 'TIP5PBOX',
    'spce': 'SPCBOX',
}

# Map the OpenMM-style name for solvent to the tleap
# name for a list of files which would enable the
# solvent model to work. Servers as error checking,
# but is not foolproof

_OPENMM_LEAP_SOLVENT_FILES_MAP = {
    'tip3p': 'leaprc.water.tip3p',
    'tip3pfb': 'leaprc.water.tip3p',
    'tip4pew': 'leaprc.water.tip4pew',
    'tip5p': 'leaprc.water.tip4pew',  # Enables the EP atom type
    'spce': 'leaprc.water.spce',
}


def remove_overlap(mol_positions, *args, **kwargs):
    """Remove any eventual overlap between a molecule and a set of others.

    The method both randomly shifts and rotates the molecule (when overlapping atoms
    are detected) until it does not clash with any other given molecule anymore. All
    the others are kept fixed.

    All the positions must be expressed in the same unit of measure.

    Parameters
    ----------
    mol_positions : numpy.ndarray
        An Nx3 array where, N is the number of atoms, containing the positions of
        the atoms of the molecule that we want to not clash with the others.
    min_distance : float
        The minimum distance accepted to consider the molecule not clashing with
        the others. Must be in the same unit of measure of the positions.
    sigma : float
        The maximum displacement for a single step. Must be in the same unit of
        measure of the positions.

    Other parameters
    ----------------
    args
        A series of numpy.ndarrays containing the positions of the atoms of the
        molecules that are kept fixed.

    Returns
    -------
    x : numpy.ndarray
        Positions of the atoms of the given molecules that do not clash.

    """
    x = np.copy(mol_positions)
    sigma = kwargs.get('sigma', 1.0)
    min_distance = kwargs.get('min_distance', 1.0)

    # Try until we have a non-overlapping conformation w.r.t. all fixed molecules
    while compute_min_dist(x, *args) <= min_distance:
        # Compute center of geometry
        x0 = x.mean(0)

        # Randomize orientation of ligand.
        Rq = mmtools.mcmc.MCRotationMove.generate_random_rotation_matrix()
        x = ((Rq * np.matrix(x - x0).T).T + x0).A

        # Choose a random displacement vector and translate
        x += sigma * np.random.randn(3)

    return x


def pack_transformation(mol1_pos, mol2_pos, min_distance, max_distance):
    """Compute an affine transformation that solve clashes and fit mol2 in the box.

    The method randomly shifts and rotates mol2 until all its atoms are within
    min_distance and max_distance from mol1. The position of mol1 is kept fixed.
    Every 200 failed iterations, the algorithm increases max_distance by 50%. It
    raise an exception after 1000 iterations.

    All the positions must be expressed in the same unit of measure.

    Parameters
    ----------
    mol1_pos : numpy.ndarray
        An Nx3 array where, N is the number of atoms, containing the positions of
        the atoms of the molecule that will be kept fixed.
    mol2_pos : numpy.ndarray
        An Nx3 array where, N is the number of atoms, containing the positions of
        the atoms of the molecule that will be eventually moved.
    min_distance : float
        The minimum distance accepted to consider mol2 not clashing with mol1. It
        must be in the same unit of measure of the positions.
    max_distance : float
        The maximum distance from mol1 to consider mol2 within the box. It must
        be in the same unit of measure of the positions.

    Returns
    -------
    transformation : numpy.ndarray
        A 4x4 ndarray representing the affine transformation that translate and
        rotate mol2.

    """
    translation = None  # we'll use this to check if we made changes to mol2_pos
    transformation = np.identity(4)

    # Compute center of geometry
    x0 = mol2_pos.mean(0)

    # Try until we have a non-overlapping conformation w.r.t. all fixed molecules
    i = 0
    min_dist, max_dist = compute_min_max_dist(mol1_pos, mol2_pos)
    while min_dist < min_distance or max_distance <= max_dist:
        # Select random atom of fixed molecule and use it to propose new x0 position
        mol1_atom_idx = np.random.random_integers(0, len(mol1_pos) - 1)
        translation = mol1_pos[mol1_atom_idx] + max_distance * np.random.randn(3) - x0

        # Generate random rotation matrix
        Rq = mmtools.mcmc.MCRotationMove.generate_random_rotation_matrix()

        # Apply random transformation and test
        x = ((Rq * np.matrix(mol2_pos - x0).T).T + x0).A + translation
        min_dist, max_dist = compute_min_max_dist(mol1_pos, x)

        # Check n iterations
        i += 1
        if i % 200 == 0:
            max_distance *= 1.5
        if i >= 1000:
            err_msg = 'Cannot fit mol2 into solvation box!'
            logger.error(err_msg)
            raise RuntimeError(err_msg)

    # Generate 4x4 affine transformation in molecule reference frame
    if translation is not None:
        transl_to_origin, transl_to_x0, rot_transl_matrix = (np.identity(4) for _ in range(3))
        transl_to_origin[:3, 3] = -x0  # translate the molecule from x0 to origin
        rot_transl_matrix[:3, :3] = Rq  # rotate molecule in origin
        rot_transl_matrix[:3, 3] = translation  # translate molecule
        transl_to_x0[:3, 3] = x0  # translate the molecule from origin to x0
        transformation = transl_to_x0.dot(rot_transl_matrix.dot(transl_to_origin))

    return transformation


def pull_close(fixed_mol_pos, translated_mol_pos, min_bound, max_bound):
    """Heuristic algorithm to quickly translate the ligand close to the receptor.

    The distance of the ligand from the receptor here is defined as the shortest
    Euclidean distance between an atom of the ligand and one of the receptor.
    The molecules positions will not be modified if the ligand is already at a
    distance in the interval [min_bound, max_bound].

    Parameters
    ----------
    fixed_mol_pos : numpy.array
        The positions of the molecule to keep fixed as a Nx3 array.
    translated_mol_pos : numpy.array
        The positions of the molecule to translate as a Nx3 array.
    min_bound : float
        Minimum distance from the receptor to the ligand. This should be high
        enough for the ligand to not overlap the receptor atoms at the beginning
        of the simulation.
    max_bound : float
        Maximum distance from the receptor to the ligand. This should be short
        enough to make the ligand and the receptor interact since the beginning
        of the simulation.

    Returns
    -------
    translation : numpy.array
        A 1x3 array containing the translation vector to apply to translated_mol_pos
        to move the molecule at a distance between min_bound and max_bound from
        fixed_mol_pos.

    """

    goal_distance = (min_bound + max_bound) / 2
    trans_pos = copy.deepcopy(translated_mol_pos)  # positions that we can modify

    # Find translation
    final_translation = np.zeros(3)
    while True:

        # Compute squared distances between all atoms
        # Each row is an array of distances from a translated atom to all fixed atoms
        # We don't need to apply square root to everything
        distances2 = np.array([((fixed_mol_pos - pos)**2).sum(1) for pos in trans_pos])

        # Find closest atoms and their distance
        min_idx = np.unravel_index(distances2.argmin(), distances2.shape)
        min_dist = np.sqrt(distances2[min_idx])

        # If closest atom is between boundaries translate ligand
        if min_bound <= min_dist <= max_bound:
            break

        # Compute unit vector that connects receptor and ligand atom
        if min_dist != 0:
            direction = fixed_mol_pos[min_idx[1]] - trans_pos[min_idx[0]]
        else:  # any deterministic direction
            direction = np.array([1, 1, 1])
        direction = direction / np.sqrt((direction**2).sum())  # normalize

        if max_bound < min_dist:  # the atom is far away
            translation = (min_dist - goal_distance) * direction
            trans_pos += translation
            final_translation += translation
        elif min_dist < min_bound:  # the two molecules overlap
            max_dist = np.sqrt(distances2.max())
            translation = (max_dist + goal_distance) * direction
            trans_pos += translation
            final_translation += translation

    return final_translation


def strip_protons(input_file_path, output_file_path):
    """
    Remove all hydrogens from PDB file and save the result.

    Input and output file cannot be the same file

    Parameters
    ----------
    input_file_path : str
        Full file path to the file to read, including extensions
    output_file_path : str
        Full file path to the file to save, including extensions
    """
    output_file = open(output_file_path, 'w')
    with open(input_file_path, 'r') as input_file:
        for line in input_file:
            if not (line[:6] == 'ATOM  ' and (line[12] == 'H' or line[13] == 'H')):
                output_file.write(line)
    output_file.close()

# For mutate_protein
_three_letter_code = {
    'A': 'ALA',
    'C': 'CYS',
    'D': 'ASP',
    'E': 'GLU',
    'F': 'PHE',
    'G': 'GLY',
    'H': 'HIS',
    'I': 'ILE',
    'K': 'LYS',
    'L': 'LEU',
    'M': 'MET',
    'N': 'ASN',
    'P': 'PRO',
    'Q': 'GLN',
    'R': 'ARG',
    'S': 'SER',
    'T': 'THR',
    'V': 'VAL',
    'W': 'TRP',
    'Y': 'TYR'
}

_one_letter_code = dict()
for one_letter in _three_letter_code.keys():
    three_letter = _three_letter_code[one_letter]
    _one_letter_code[three_letter] = one_letter


def decompose_mutation(mutation):
    match = re.match('(\D)(\d+)(\D)', mutation)
    try:
        original_residue_name = _three_letter_code[match.group(1)]
        residue_index = int(match.group(2))
        mutated_residue_name = _three_letter_code[match.group(3)]
    except AttributeError:
        error_string = 'Mutation "{}" could not be parsed! '.format(mutation)
        error_string += 'Should be of form {single letter}{integer}{another single letter}'
        raise ValueError(error_string)
    return original_residue_name, residue_index, mutated_residue_name


def generate_pdbfixer_mutation_code(original_residue_name, residue_index, mutated_residue_name):
    return '{0:s}-{1:d}-{2:s}'.format(original_residue_name, residue_index, mutated_residue_name)


def process_tool_directive(directives, option, dispatch, allowed_values, yields_value=False):
    """Process a directive.

    Parameters
    ----------
    option : str
        The name of the option to be processed.
        Will remove this option from `directives` once processed.
    dispatch : function
        The function to call.
    allowed_values : list
        If not None, the value of directives[option] will be checked against this list
    yields_value : boolean, default False
        Tells this function to expect a return from the dispatch function and give it back as needed
    """
    if option in directives:
        value = directives[option]
        # Validate options
        if allowed_values is not None:
            if value not in allowed_values:
                raise ValueError("'{}' must be one of {}".format(option, allowed_values))
        # Dispatch
        output = dispatch(value)
        # Delete the key once we've processed it
        del directives[option]
        if yields_value:
            return output
        return

def apply_pdbfixer(input_file_path, output_file_path, directives):
    """
    Apply PDBFixer to make changes to the specified molecule.

    Single mutants are supported in the form "T315I"
    Double mutants are supported in the form "L858R/T790M"

    The string "WT" still pushes the molecule through PDBFixer, but makes no mutations.
    This is useful for testing.

    Original PDB file numbering scheme is used.

    Currently, only PDB files are supported.
    pdbfixer is used to make the mutations

    Parameters
    ----------
    input_file_path : str
        Full file path to the file to read, including extensions
    output_file_path : str
        Full file path to the file to save, including extensions
    directives : dict
        Dict containing directives for PDBFixer.
    """
    DEFAULT_PH = 7.4 # default pH

    # Make a copy since we will delete from the dictionary to validate
    directives = copy.deepcopy(directives)

    # Create a PDBFixer object
    fixer = PDBFixer(input_file_path)
    fixer.missingResidues = {}

    # Dispatch functions
    # These won't be documented individually because they are so short
    def dispatch_pH(value):
        pH = DEFAULT_PH
        try:
            pH = float(value)
            logger.info('pdbfixer: Will use user-specified pH {}'.format(pH))
        except:
            raise ValueError("'ph' must be a floating-point number: found '{}'".format(value))
        return pH

    pH = process_tool_directive(directives, 'ph', dispatch_pH, None, yields_value=True)

    def add_missing_residues(value):
        if value == 'yes':
            fixer.findMissingResidues()
            logger.info('pdbfixer: Will add missing residues specified in SEQRES')

    def apply_mutations(value):
        # Extract chain id
        chain_id = None
        if 'chain_id' in value:
            chain_id = value['chain_id']
            if chain_id == 'none':
                chain_id = None
        # Extract mutations
        mutations = value['mutations']
        # Convert mutations to PDBFixer format
        if mutations != 'WT':
            pdbfixer_mutations = [generate_pdbfixer_mutation_code(*decompose_mutation(mutation))
                                  for mutation in mutations.split('/')]
            logger.info('pdbfixer: Will make mutations {} to chain_id {}.'.format(pdbfixer_mutations, chain_id))

            fixer.applyMutations(pdbfixer_mutations, chain_id)
        else:
            logger.info('pdbfixer: No mutations will be applied since "WT" specified.')

    def replace_nonstandard_residues(value):
        if value == 'yes':
            logger.info('pdbfixer: Will replace nonstandard residues.')
            fixer.findNonstandardResidues()
            fixer.replaceNonstandardResidues()

    def remove_heterogens(value):
        if value == 'water':
            logger.info('pdbfixer: Will remove heterogens, retaining water.')
            fixer.removeHeterogens(keepWater=True)
        elif value == 'all':
            logger.info('pdbfixer: Will remove heterogens, discarding water.')
            fixer.removeHeterogens(keepWater=False)

    def add_missing_atoms(value):
        fixer.findMissingAtoms()
        if value not in ('all', 'heavy'):
            fixer.missingAtoms = {}
            fixer.missingTerminals = {}
        logger.info('pdbfixer: Will add missing atoms: {}.'.format(value))
        fixer.addMissingAtoms()
        if value in ('all', 'hydrogens'):
            logger.info('pdbfixer: Will add hydrogens in default protonation state for pH {}.'.format(pH))
            fixer.addMissingHydrogens(pH)

    # Set default atom addition method
    if 'add_missing_atoms' not in directives:
        directives['add_missing_atoms'] = 'heavy'

    # Dispatch directives
    process_tool_directive(directives, 'add_missing_residues', add_missing_residues, [True, False])
    process_tool_directive(directives, 'apply_mutations', apply_mutations, None)
    process_tool_directive(directives, 'replace_nonstandard_residues', replace_nonstandard_residues, [True, False])
    process_tool_directive(directives, 'remove_heterogens', remove_heterogens, ['all', 'water', 'none'])
    process_tool_directive(directives, 'add_missing_atoms', add_missing_atoms, ['all', 'heavy', 'hydrogens', 'none'])

    # Check that there were no extra options
    if len(directives) > 0:
        raise ValueError("The 'pdbfixer:' block contained some nodes that it didn't know how to process: {}".format(directives))

    # Write the final structure
    PDBFile.writeFile(fixer.topology, fixer.positions, open(output_file_path, 'w'))


def apply_modeller(input_file_path, output_file_path, directives):
    """
    Apply Salilab Modeller to make changes to the specified molecule.

    Single mutants are supported in the form "T315I"
    Double mutants are not currently supported.

    The string "WT" makes no mutations.

    Original PDB file numbering scheme is used.

    Currently, only PDB files are supported.
    modeller is used to make the mutations. You must have a license file installed for this to work

    Parameters
    ----------
    input_file_path : str
        Full file path to the file to read, including extensions
    output_file_path : str
        Full file path to the file to save, including extensions
    directives : dict
        Dict containing directives for modeller.
    """

    if not utils.is_modeller_installed():
        raise ImportError('Modeller and license must be installed to use this feature.')
    import modeller

    directives = copy.deepcopy(directives)

    # Silence unnecessary output to the log files
    modeller.log.none()

    # Create modeller environment and point it to the PDB file
    env = modeller.environ()
    atom_files_directory = os.path.dirname(input_file_path)
    atom_file_name = os.path.basename(input_file_path)

    # Read in topology and parameter files
    env.libs.topology.read(file='$(LIB)/top_heav.lib')
    env.libs.parameters.read(file='$(LIB)/par.lib')

    env.io.atom_files_directory = [atom_files_directory]
    alignment = modeller.alignment(env)
    model = modeller.model(env, file=atom_file_name)
    model_original_numbering = modeller.model(env, file=atom_file_name)
    alignment.append_model(model, atom_files=atom_file_name, align_codes=atom_file_name)

    def apply_mutations_modeller(value):
        # Extract chain id
        chain_id = None
        if 'chain_id' in value:
            chain_id = value['chain_id']
            if chain_id == 'none':
                chain_id = 0
        # Extract mutations
        mutations = value['mutations']
        # Convert mutations to PDBFixer format
        if mutations != 'WT':
            modeller_mutations = [generate_pdbfixer_mutation_code(*decompose_mutation(mutation))
                                  for mutation in mutations.split('/')]
            if len(modeller_mutations) > 1:
                raise ValueError('{} is a double mutant and not supported by Modeller currently.'.format(mutations))
            else:
                logger.info('modeller: Will make mutations {} to chain_id {}.'.format(modeller_mutations, chain_id))
                sel = modeller.selection(model.chains[chain_id].residues[modeller_mutations[0].split('-')[1]])
                sel.mutate(residue_type=modeller_mutations[0].split('-')[2])
                alignment.append_model(model, align_codes=modeller_mutations[0])
                model.clear_topology()
                model.generate_topology(alignment[modeller_mutations[0]])
                model.transfer_xyz(alignment)
                model.build(initialize_xyz=False, build_method='INTERNAL_COORDINATES')

        else:
            logger.info('modeller: No mutations will be applied since "WT" specified.')
            alignment.append_model(model, align_codes='WT')

    process_tool_directive(directives, 'apply_mutations', apply_mutations_modeller, None)

    # Check that there were no extra options
    if len(directives) > 0:
        raise ValueError("The 'modeller:' block contained some nodes that it didn't know how to process: {}".format(directives))

    # Write the final model
    model.res_num_from(model_original_numbering, alignment)
    model.write(file=output_file_path)


def read_csv_lines(file_path, lines):
    """Return a list of CSV records.

    The function takes care of ignoring comments and blank lines.

    Parameters
    ----------
    file_path : str
        The path to the CSV file.
    lines : 'all' or int
        The index of the line to read or 'all' to return
        the list of all lines.

    Returns
    -------
    records : str or list of str
        The CSV record if lines is an integer, or a list of CSV
        records if it is 'all'.

    """
    # Read all lines ignoring blank lines and comments.
    with open(file_path, 'r') as f:
        all_records = [line for line in f
                       if bool(line) and not line.strip().startswith('#')]

    if lines == 'all':
        return all_records
    return all_records[lines]


# ==============================================================================
# SETUP DATABASE
# ==============================================================================

class SetupDatabase:
    """Provide utility functions to set up systems and molecules.

    The object allows to access molecules, systems and solvents by 'id' and takes
    care of parametrizing molecules and creating the AMBER prmtop and inpcrd files
    describing systems.

    Parameters
    ----------
    setup_dir : str
        Path to the main setup directory. Changing this means changing the database.
    molecules : dict, Optional. Default: None
        YAML description of the molecules.
        Dictionary should be of form {molecule_id : molecule YAML description}
    solvents : dict, Optional. Default: None
        YAML description of the solvents.
        Dictionary should be of form {solvent_id : solvent YAML description}
    systems : dict, Optional. Default: None
        YAML description of the systems.
        Dictionary should be of form {system_id : system YAML description}

    """

    SYSTEMS_DIR = 'systems'  #: Stock system's sub-directory name
    MOLECULES_DIR = 'molecules'  #: Stock Molecules sub-directory name
    CLASH_THRESHOLD = 1.5  #: distance in Angstroms to consider two atoms clashing

    def __init__(self, setup_dir, molecules=None, solvents=None, systems=None):
        """Initialize the database."""
        self.setup_dir = setup_dir
        self.molecules = molecules
        self.solvents = solvents
        self.systems = systems

        # Private attributes
        self._pos_cache = {}  # cache positions of molecules
        self._processed_mols = set()  # keep track of parametrized molecules

    def get_molecule_dir(self, molecule_id):
        """Return the directory where the parameter files are stored.

        Parameters
        ----------
        molecule_id : str
            The ID of the molecule.

        Returns
        -------
        str
            The path to the molecule directory.
        """
        return os.path.join(self.setup_dir, self.MOLECULES_DIR, molecule_id)

    def get_system_files_paths(self, system_id):
        """Return the paths to the systems files.

        Parameters
        ----------
        system_id : str
            The ID of the system.

        Returns
        -------
        system_files_paths : list of namedtuple
            Elements of the list contain the paths to the system files for
            each phase. Each namedtuple contains the fields position_path (e.g.
            inpcrd, gro, or pdb) and parameters_path (e.g. prmtop, top, or xml).

        """
        Paths = collections.namedtuple('Paths', ['position_path', 'parameters_path'])
        system_dir = os.path.join(self.setup_dir, self.SYSTEMS_DIR, system_id)

        if 'receptor' in self.systems[system_id]:
            system_files_paths = [
                Paths(position_path=os.path.join(system_dir, 'complex.inpcrd'),
                      parameters_path=os.path.join(system_dir, 'complex.prmtop')),
                Paths(position_path=os.path.join(system_dir, 'solvent.inpcrd'),
                      parameters_path=os.path.join(system_dir, 'solvent.prmtop'))
            ]
        elif 'solute' in self.systems[system_id]:
            system_files_paths = [
                Paths(position_path=os.path.join(system_dir, 'solvent1.inpcrd'),
                      parameters_path=os.path.join(system_dir, 'solvent1.prmtop')),
                Paths(position_path=os.path.join(system_dir, 'solvent2.inpcrd'),
                      parameters_path=os.path.join(system_dir, 'solvent2.prmtop'))
            ]
        else:
            parameter_file_extensions = {'prmtop', 'top', 'xml', 'psf'}
            system_files_paths = []
            for phase_path_name in ['phase1_path', 'phase2_path']:
                file_paths = self.systems[system_id][phase_path_name]
                assert len(file_paths) == 2

                # Make sure that the position file is first.
                first_file_extension = os.path.splitext(file_paths[0])[1][1:]
                if first_file_extension in parameter_file_extensions:
                    file_paths = list(reversed(file_paths))

                # Append Paths object.
                system_files_paths.append(Paths(position_path=file_paths[0],
                                                parameters_path=file_paths[1]))
        return system_files_paths

    def is_molecule_setup(self, molecule_id):
        """Check whether the molecule has been processed previously.

        The molecule must be set up if it needs to be parametrize by antechamber
        (and the gaff.mol2 and frcmod files do not exist), if the molecule must be
        generated by OpenEye, or if it needs to be extracted by a multi-molecule file.

        An example to clarify the difference between the two return values: a protein
        in a single-frame pdb does not have to be processed (since it does not go through
        antechamber) thus the function will return ``is_setup=True`` and ``is_processed=False``.

        Parameters
        ----------
        molecule_id : str
            The id of the molecule.

        Returns
        -------
        is_setup : bool
            True if the molecule's parameter files have been specified by the user
            or if they have been generated by SetupDatabase.
        is_processed : bool
            True if parameter files have been generated previously by SetupDatabase
            (i.e. if the parameter files were not manually specified by the user).

        """

        # The only way to check if we processed the molecule in the current run is
        # through self._processed_mols as 'parameters' will be changed after setup
        if molecule_id in self._processed_mols:
            return True, True

        # Some convenience variables
        molecule_descr = self.molecules[molecule_id]
        molecule_dir = self.get_molecule_dir(molecule_id)
        molecule_id_path = os.path.join(molecule_dir, molecule_id)
        try:
            extension = os.path.splitext(molecule_descr['filepath'])[1]
        except KeyError:
            extension = None

        # The following checks must be performed in reverse order w.r.t. how they
        # are executed in _setup_molecules()

        files_to_check = {}

        # If the molecule must go through antechamber we search for its output
        if 'antechamber' in molecule_descr:
            files_to_check = [('filepath', molecule_id_path + '.gaff.mol2'),
                              (['leap', 'parameters'], molecule_id_path + '.frcmod')]

        # If the molecule must be generated by OpenEye, a mol2 should have been created
        elif extension is None or extension == '.smiles' or extension == '.csv':
            files_to_check = [('filepath', molecule_id_path + '.mol2')]

        # If we have to strip the protons off a PDB, a new PDB should have been created
        elif 'strip_protons' in molecule_descr and molecule_descr['strip_protons']:
            files_to_check = [('filepath', molecule_id_path + '.pdb')]

        # If we have to make mutations, a new PDB should be created
        elif 'pdbfixer' in molecule_descr:
            files_to_check = [('filepath', molecule_id_path + '.pdb')]

        # If we have to make mutations using modeller, a new PDB should be created
        elif 'modeller' in molecule_descr:
            files_to_check = [('filepath', molecule_id_path + '.pdb')]

        # If a single structure must be extracted we search for output
        elif 'select' in molecule_descr:
            files_to_check = [('filepath', molecule_id_path + extension)]

        # Check if this needed to be processed at all
        if not files_to_check:
            return True, False

        # Check if all output files exist
        all_file_exist = True
        for descr_key, file_path in files_to_check:
            all_file_exist &= os.path.isfile(file_path) and os.path.getsize(file_path) > 0
            if all_file_exist:  # Make sure internal description is correct
                try:
                    molecule_descr[descr_key] = file_path
                except TypeError:  # nested key, list are unhashable
                    molecule_descr[descr_key[0]][descr_key[1]].append(file_path)

        # Compute and update small molecule net charge
        if all_file_exist:
            extension = os.path.splitext(molecule_descr['filepath'])[1]
            # TODO what if this is a peptide? This should be computed in get_system()
            if extension == '.mol2':
                molecule_descr['net_charge'] = utils.Mol2File(molecule_descr['filepath']).net_charge

        return all_file_exist, all_file_exist

    def is_system_setup(self, system_id):
        """Check whether the system has been already processed.

        Parameters
        ----------
        system_id : str
            The ID of the system.

        Returns
        -------
        is_setup : bool
            True if the system is ready to be used for an experiment. Either because
            the system has directly provided the system files, or because it already
            went through the setup pipeline.
        is_processed : bool
            True if the system has already gone through the setup pipeline.

        """
        if 'ligand' in self.systems[system_id] or 'solute' in self.systems[system_id]:
            system_files_paths = self.get_system_files_paths(system_id)
            is_setup = (os.path.exists(system_files_paths[0].position_path) and
                        os.path.exists(system_files_paths[0].parameters_path) and
                        os.path.exists(system_files_paths[1].position_path) and
                        os.path.exists(system_files_paths[1].parameters_path))
            return is_setup, is_setup
        else:
            return True, False

    def get_system(self, system_id):
        """Make sure that the system files are set up and return the system folder.

        If necessary, create the prmtop and inpcrd files from the given components.
        The system files are generated with tleap. If no molecule specifies a general
        force field, leaprc.ff14SB is loaded.

        Parameters
        ----------
        system_id : str
            The ID of the system.

        Returns
        -------
        system_files_paths : list of namedtuple
            Elements of the list contain the paths to the system files for
            each phase. Each namedtuple contains the fields position_path (e.g.
            inpcrd, gro, or pdb) and parameters_path (e.g. prmtop, top, or xml).

        """
        # Check if system has been already processed
        system_files_paths = self.get_system_files_paths(system_id)
        if self.is_system_setup(system_id)[0]:
            return system_files_paths

        system_descr = self.systems[system_id]
        log_message = 'Setting up the systems for {} and {} using solvent {}'

        if 'receptor' in system_descr:  # binding free energy calculation
            receptor_id = system_descr['receptor']
            ligand_id = system_descr['ligand']
            solvent_id = system_descr['solvent']
            system_parameters = system_descr['leap']['parameters']
            logger.info(log_message.format(receptor_id, ligand_id, solvent_id))

            # solvent phase
            logger.debug('Setting up solvent phase')
            self._setup_system(system_files_paths[1].position_path, False,
                               0, system_parameters, solvent_id, ligand_id)

            try:
                alchemical_charge = int(round(self.molecules[ligand_id]['net_charge']))
            except KeyError:
                alchemical_charge = 0

            # complex phase
            logger.debug('Setting up complex phase')
            self._setup_system(system_files_paths[0].position_path,
                               system_descr['pack'], alchemical_charge,
                               system_parameters, solvent_id, receptor_id,
                               ligand_id)
        else:  # partition/solvation free energy calculation
            solute_id = system_descr['solute']
            solvent1_id = system_descr['solvent1']
            solvent2_id = system_descr['solvent2']
            system_parameters = system_descr['leap']['parameters']
            logger.info(log_message.format(solute_id, solvent1_id, solvent2_id))

            # solvent1 phase
            logger.debug('Setting up solvent1 phase')
            self._setup_system(system_files_paths[0].position_path, False,
                               0, system_parameters, solvent1_id, solute_id)

            # solvent2 phase
            logger.debug('Setting up solvent2 phase')
            self._setup_system(system_files_paths[1].position_path, False,
                               0, system_parameters, solvent2_id, solute_id)

        return system_files_paths

    def setup_all_systems(self):
        """Setup all molecules and systems in the database.

        The method supports parallelization through MPI.

        """
        # Find all molecules that need to be set up.
        molecules_to_setup = []
        for molecule_id in self.molecules:
            if not self.is_molecule_setup(molecule_id)[0]:
                molecules_to_setup.append(molecule_id)
        molecules_to_setup.sort()

        # Parallelize generation of all molecules among nodes.
        mpiplus.distribute(self._setup_molecules,
                           distributed_args=molecules_to_setup,
                           send_results_to=None, group_size=1, sync_nodes=True)

        # Find all systems that need to be set up.
        systems_to_setup = []
        for system_id in self.systems:
            if not self.is_system_setup(system_id)[0]:
                systems_to_setup.append(system_id)
        systems_to_setup.sort()

        # Parallelize generation of all systems among nodes.
        mpiplus.distribute(self.get_system,
                           distributed_args=systems_to_setup,
                           send_results_to=None, group_size=1, sync_nodes=True)

    def _generate_molecule(self, molecule_id):
        """Generate molecule using the OpenEye toolkit from name or smiles.

        The molecule is charged with OpenEye's recommended AM1BCC charge
        selection scheme and it is saved into the OpenEye molecules cache.

        Parameters
        ----------
        molecule_id : str
            The id of the molecule as given in the YAML script

        Returns
        -------
        molecule : OEMol
            The generated molecule.

        """
        mol_descr = self.molecules[molecule_id]  # molecule description
        try:
            if 'name' in mol_descr:
                molecule = moltools.openeye.iupac_to_oemol(mol_descr['name'])
            elif 'smiles' in mol_descr:
                molecule = moltools.openeye.smiles_to_oemol(mol_descr['smiles'])
            molecule = moltools.openeye.get_charges(molecule, keep_confs=1)
        except ImportError as e:
            error_msg = ('requested molecule generation from name or smiles but '
                         'could not find OpenEye toolkit: ' + str(e))
            logger.error(error_msg)
            raise RuntimeError(error_msg)

        return molecule

    def _generate_residue_name(self, molecule_id):
        """Generates a residue name for a molecule.

        The function guarantees to not generate twice the same residue name.
        Purely numeric residue names mess up the pipeline, so we generate
        residue names of the form YXX, where Y is a letter and X are digits
        (e.g. A01, A02, ..., Z99).

        The order of the generated residue name is not guaranteed.

        WARNING: The algorithm may fail when new molecules are added to
        self.molecules after construction. This is not the case right now,
        but it's good to keep in mind.

        Parameters
        ----------
        molecule_id : str
            The molecule identifier.

        Returns
        -------
        residue_name : str
            A three-character residue name.
        """
        # We need to associate a unique number to this molecule, and do
        # so in such a way that distributing molecule setups over multiple
        # MPI process still ends up in unique residue names for each molecule.
        molecule_ids = sorted(self.molecules.keys())
        n_molecule = molecule_ids.index(molecule_id)
        assert n_molecule < 2600

        # Build 3-character identifier.
        character = chr(n_molecule // 100 + 65)
        digits = str(n_molecule % 100)
        residue_name = character + digits.zfill(2)
        return residue_name

    def _setup_molecules(self, *args):
        """Set up the files needed to generate the system for all the molecules.

        If OpenEye tools are installed, this generate the molecules when the source is
        not a file. If two (or more) molecules generated by OpenEye have overlapping
        atoms, the molecules are randomly shifted and rotated until the clash is resolved.
        With the OpenEye toolkit installed, we also perform a sanity check to verify that
        the molecules from files do not have overlapping atoms. An Exception is raised if
        this is not the case.

        If the Schrodinger's suite is install, this can enumerate tautomeric and protonation
        states with epik when requested.

        This also parametrize the molecule with antechamber when requested.

        Other parameters
        ----------------
        args
            All the molecules ids that compose the system. These molecules are the only
            ones considered when trying to resolve the overlapping atoms.

        """

        for mol_id in args:
            net_charge = None  # used by antechamber
            mol_descr = self.molecules[mol_id]

            # Have we already processed this molecule? Do we have to do it at all?
            # We don't want to create the output folder if we don't need to
            if self.is_molecule_setup(mol_id)[0]:
                continue

            # Create output directory if it doesn't exist
            mol_dir = self.get_molecule_dir(mol_id)
            if not os.path.exists(mol_dir):
                os.makedirs(mol_dir)

            try:
                extension = os.path.splitext(mol_descr['filepath'])[1]
            except KeyError:
                extension = None

            # Extract single model if this is a multi-model file
            if 'select' in mol_descr:
                model_idx = mol_descr['select']
                single_file_path = os.path.join(mol_dir, mol_id + extension)
                if extension == '.pdb':
                    # Create single-model PDB file
                    pdb_file = PDBFile(mol_descr['filepath'])
                    with open(single_file_path, 'w') as f:
                        PDBFile.writeHeader(pdb_file.topology, file=f)
                        PDBFile.writeModel(pdb_file.topology, pdb_file.getPositions(frame=model_idx), file=f)
                    # We might as well already cache the positions
                    self._pos_cache[mol_id] = pdb_file.getPositions(asNumpy=True, frame=model_idx) / unit.angstrom

                elif extension == '.smiles' or extension == '.csv':
                    # Extract the correct line and save it in a new file.
                    smiles_line = read_csv_lines(mol_descr['filepath'], lines=model_idx)
                    with open(single_file_path, 'w') as f:
                        f.write(smiles_line)

                elif extension == '.mol2' or extension == '.sdf':
                    if not utils.is_openeye_installed(oetools=('oechem',)):
                        raise RuntimeError('Cannot support {} files selection without OpenEye'.format(
                                extension[1:]))
                    oe_molecule = utils.load_oe_molecules(mol_descr['filepath'], molecule_idx=model_idx)
                    if extension == '.mol2':
                        mol_names = list(utils.Mol2File(mol_descr['filepath']).resnames)
                        utils.write_oe_molecule(oe_molecule, single_file_path, mol2_resname=mol_names[model_idx])
                    else:
                        utils.write_oe_molecule(oe_molecule, single_file_path)

                else:
                    raise RuntimeError('Model selection is not supported for {} files'.format(extension[1:]))

                # Save new file path
                mol_descr['filepath'] = single_file_path

            # Strip off protons if required
            if 'strip_protons' in mol_descr and mol_descr['strip_protons']:
                if extension != '.pdb':
                    raise RuntimeError('Cannot strip protons from {} files.'.format(extension[1:]))
                output_file_path = os.path.join(mol_dir, mol_id + '.pdb')
                strip_protons(mol_descr['filepath'], output_file_path)
                mol_descr['filepath'] = output_file_path

            # Apply PDBFixer if requested
            if 'pdbfixer' in mol_descr:
                if extension not in ['.pdb', '.PDB']:
                    raise RuntimeError('Cannot apply PDBFixer to {} files; a .pdb file is required.'.format(extension[1:]))
                output_file_path = os.path.join(mol_dir, mol_id + '.pdb')
                apply_pdbfixer(mol_descr['filepath'], output_file_path, mol_descr['pdbfixer'])
                mol_descr['filepath'] = output_file_path

            # Apply modeller if requested
            if 'modeller' in mol_descr:
                if extension not in ['.pdb', '.PDB']:
                    raise RuntimeError('Cannot apply modeller to {} files; a .pdb file is required.'.format(extension[1:]))
                output_file_path = os.path.join(mol_dir, mol_id + '.pdb')
                apply_modeller(mol_descr['filepath'], output_file_path, mol_descr['modeller'])
                mol_descr['filepath'] = output_file_path

            # Generate missing molecules with OpenEye. At the end of parametrization
            # we update the 'filepath' key also for OpenEye-generated molecules so
            # we don't need to keep track of the molecules we have already generated
            if extension is None or extension == '.smiles' or extension == '.csv':
                if not utils.is_openeye_installed(oetools=('oechem', 'oeiupac', 'oequacpac', 'oeomega')):
                    if extension is None:
                        raise RuntimeError('Cannot generate molecule {} without OpenEye licensed with '
                                           'OEChem, OEIUPAC, OEOmega, and OEQuacPack.'.format(mol_id))
                    else:
                        raise RuntimeError('Cannot support {} files without OpenEye licensed with '
                                           'OEChem, OEIUPAC, OEOmega, and OEQuacPack.'.format(extension[1:]))

                # Retrieve the first SMILES string (eventually extracted
                # while handling of the 'select' keyword above)
                if extension is not None:
                    # Get first record in CSV file.
                    first_line = read_csv_lines(mol_descr['filepath'], lines=0)

                    # Automatically detect if delimiter is comma or semicolon
                    for delimiter in ',;':
                        logger.debug("Attempt to parse smiles file with delimiter '{}'".format(delimiter))
                        line_fields = first_line.split(delimiter)
                        # If there is only one column, take that, otherwise take second
                        if len(line_fields) > 1:
                            smiles_str = line_fields[1].strip()
                        else:
                            smiles_str = line_fields[0].strip()

                        # try to generate the smiles and try new delimiter if it fails
                        mol_descr['smiles'] = smiles_str
                        try:
                            oe_molecule = self._generate_molecule(mol_id)
                            break
                        except (ValueError, RuntimeError):
                            oe_molecule = None

                    # Raise an error if no delimiter worked
                    if oe_molecule is None:
                        raise RuntimeError('Cannot detect SMILES file format.')
                else:
                    # Generate molecule from mol_descr['smiles']
                    oe_molecule = self._generate_molecule(mol_id)

                # Cache atom positions
                self._pos_cache[mol_id] = utils.get_oe_mol_positions(oe_molecule)

                # Write OpenEye generated molecules in mol2 files
                # We update the 'filepath' key in the molecule description
                mol_descr['filepath'] = os.path.join(mol_dir, mol_id + '.mol2')

                # Generate a residue name for the SMILES molecule.
                residue_name = self._generate_residue_name(mol_id)
                moltools.openeye.molecule_to_mol2(oe_molecule, mol_descr['filepath'],
                                                  residue_name=residue_name)

            # Enumerate protonation states with epik
            if 'epik' in mol_descr:
                epik_base_path = os.path.join(mol_dir, mol_id + '-epik.')
                epik_mae_file = epik_base_path + 'mae'
                epik_mol2_file = epik_base_path + 'mol2'
                epik_sdf_file = epik_base_path + 'sdf'

                # Run epik and convert from maestro to both mol2 and sdf
                # to not lose neither the penalties nor the residue name
                epik_kwargs = mol_descr['epik']
                moltools.schrodinger.run_epik(mol_descr['filepath'], epik_mae_file, **epik_kwargs)
                moltools.schrodinger.run_structconvert(epik_mae_file, epik_sdf_file)
                moltools.schrodinger.run_structconvert(epik_mae_file, epik_mol2_file)

                # Save new net charge from the i_epik_Tot_Q property
                net_charge = int(moltools.schrodinger.run_proplister(epik_sdf_file)[0]['i_epik_Tot_Q'])

                # Keep filepath consistent
                mol_descr['filepath'] = epik_mol2_file

            # Antechamber does not support sdf files so we need to convert them
            extension = os.path.splitext(mol_descr['filepath'])[1]
            if extension == '.sdf':
                if not utils.is_openeye_installed(oetools=('oechem',)):
                    raise RuntimeError('Cannot support sdf files without OpenEye OEChem')
                mol2_file_path = os.path.join(mol_dir, mol_id + '.mol2')
                oe_molecule = utils.load_oe_molecules(mol_descr['filepath'], molecule_idx=0)

                # Generate a residue name for the sdf molecule.
                residue_name = self._generate_residue_name(mol_id)
                moltools.openeye.molecule_to_mol2(oe_molecule, mol2_file_path,
                                                  residue_name=residue_name)

                # Update filepath information
                mol_descr['filepath'] = mol2_file_path

            # Parametrize the molecule with antechamber
            if 'antechamber' in mol_descr:
                # Generate charges with OpenEye if requested
                if 'openeye' in mol_descr:
                    if not utils.is_openeye_installed(oetools=('oechem', 'oequacpac', 'oeomega')):
                        err_msg = ('Cannot find OpenEye toolkit with OEChem and OEQuacPac to compute charges '
                                   'for molecule {}').format(mol_id)
                        logger.error(err_msg)
                        raise RuntimeError(err_msg)
                    mol2_file_path = os.path.join(mol_dir, mol_id + '.mol2')
                    oe_molecule = utils.load_oe_molecules(mol_descr['filepath'], molecule_idx=0)

                    # Setting keep_confs = None keeps the original conformation
                    oe_molecule = moltools.openeye.get_charges(oe_molecule, keep_confs=None)
                    residue_name = utils.Mol2File(mol_descr['filepath']).resname
                    moltools.openeye.molecule_to_mol2(oe_molecule, mol2_file_path,
                                                      residue_name=residue_name)

                    utils.Mol2File(mol2_file_path).round_charge()  # normalize charges
                    # We don't need Epik's or input net charge as antechamber will
                    # infer the net charge from the sum of the OpenEye charges.
                    net_charge = None
                    mol_descr['filepath'] = mol2_file_path
                # Check if use specified a net_charge, but don't overwrite Epik's protonation state.
                elif net_charge is not None:
                    net_charge = mol_descr['antechamber'].get('net_charge', None)

                # Generate parameters
                charge_method = mol_descr['antechamber']['charge_method']
                input_mol_path = os.path.abspath(mol_descr['filepath'])
                # Use Gaff in parameters, otherwise default to gaff2
                gaff = 'gaff2' if 'leaprc.gaff2' in mol_descr['leap']['parameters'] else 'gaff'
                with moltools.utils.temporary_cd(mol_dir):
                    moltools.amber.run_antechamber(mol_id, input_mol_path,
                                                   charge_method=charge_method,
                                                   net_charge=net_charge,
                                                   gaff_version=gaff)

                # Save new parameters paths
                mol_descr['filepath'] = os.path.join(mol_dir, mol_id + '.gaff.mol2')
                mol_descr['leap']['parameters'].append(os.path.join(mol_dir, mol_id + '.frcmod'))

                # Normalize charges if not done before
                if 'openeye' not in mol_descr:
                    utils.Mol2File(mol_descr['filepath']).round_charge()

            # Determine small molecule net charge
            extension = os.path.splitext(mol_descr['filepath'])[1]
            if extension == '.mol2':
                # TODO what if this is a peptide? this should be computed in get_system()
                mol_descr['net_charge'] = utils.Mol2File(mol_descr['filepath']).net_charge

            # Keep track of processed molecule
            self._processed_mols.add(mol_id)

    def _setup_system(self, system_file_path, pack, alchemical_charge,
                      system_parameters, solvent_id, *molecule_ids, **kwargs):
        """Setup a system and create its prmtop/inpcrd files.

        IMPORTANT: This function does not check if it's about to overwrite
        files. Use get_system() for safe setup.

        Parameters
        ----------
        system_file_path : str
            The path to either the prmtop or inpcrd output file. The other one
            will be saved in the same folder with the same base name.
        pack : bool
            True to automatically solve atom clashes and reduce box dimension.
        alchemical_charge : int
            Number of counterions to alchemically modify during the simulation.
        system_parameters : list of str
            Contain the parameters file that must be loaded in tleap for the
            system in addition to the molecule-specific ones.
        solvent_id : str
            The ID of the solvent.
        ignore_ionic_strength : bool, optional
            If True, no ions will be added to reach the ionic strength (default
            is False).
        save_amber_files : bool, optional
            If False, prmtop and inpcrd files are not saved (default is True).

        Other Parameters
        ----------------
        *molecule_ids : list-like of str
            List the IDs of the molecules to pack together in the system.

        """

        # Get kwargs
        ignore_ionic_strength = kwargs.pop('ignore_ionic_strength', False)
        save_amber_files = kwargs.pop('save_amber_files', True)
        assert len(kwargs) == 0

        # Make sure molecules are set up
        self._setup_molecules(*molecule_ids)
        solvent = self.solvents[solvent_id]

        # Start error tracking variables
        # Water
        known_solvent_files = [file for file in _OPENMM_LEAP_SOLVENT_FILES_MAP.values()]
        loaded_water_files = []  # Detected loaded water files

        def extend_list_of_waters(leap_parameters):
            """Extend the loaded_water_files list given the leap_parameters"""
            loaded_water_files.extend([water for water
                                       in leap_parameters
                                       if water in known_solvent_files])

        # Create tleap script
        tleap = utils.TLeap()

        # Load all parameters
        # --------------------
        tleap.new_section('Load parameters')

        for mol_id in molecule_ids:
            molecule_parameters = self.molecules[mol_id]['leap']['parameters']
            # Track loaded water models
            extend_list_of_waters(molecule_parameters)
            tleap.load_parameters(*molecule_parameters)

        extend_list_of_waters(system_parameters)
        tleap.load_parameters(*system_parameters)
        solvent_leap = solvent['leap']['parameters']
        extend_list_of_waters(solvent_leap)
        tleap.load_parameters(*solvent_leap)

        # Load molecules and create complexes
        # ------------------------------------
        tleap.new_section('Load molecules')
        for mol_id in molecule_ids:
            tleap.load_unit(unit_name=mol_id, file_path=self.molecules[mol_id]['filepath'])

        if len(molecule_ids) > 1:
            # Check that molecules don't have clashing atoms. Also, if the ligand
            # is too far away from the molecule we want to pull it closer
            # TODO this check should be available even without OpenEye
            if pack and utils.is_openeye_installed(oetools=('oechem',)):

                # Load atom positions of all molecules
                positions = [0 for _ in molecule_ids]
                for i, mol_id in enumerate(molecule_ids):
                    if mol_id not in self._pos_cache:
                        self._pos_cache[mol_id] = utils.get_oe_mol_positions(
                                utils.load_oe_molecules(self.molecules[mol_id]['filepath'], molecule_idx=0))
                    positions[i] = self._pos_cache[mol_id]

                # Find and apply the transformation to fix clashing
                # TODO this doesn't work with more than 2 molecule_ids
                try:
                    max_dist = solvent['clearance'].value_in_unit(unit.angstrom) / 1.5
                except KeyError:
                    max_dist = 10.0
                transformation = pack_transformation(positions[0], positions[1],
                                                     self.CLASH_THRESHOLD, max_dist)
                if (transformation != np.identity(4)).any():
                    logger.warning('Changing starting positions for {}.'.format(molecule_ids[1]))
                    tleap.new_section('Fix clashing atoms')
                    tleap.transform(molecule_ids[1], transformation)

            # Create complex
            tleap.new_section('Create complex')
            tleap.combine('complex', *molecule_ids)
            unit_to_solvate = 'complex'
        else:
            unit_to_solvate = molecule_ids[0]

        # Configure solvent
        # ------------------
        if solvent['nonbonded_method'] == openmm.app.NoCutoff:
            try:
                implicit_solvent = solvent['implicit_solvent']
            except KeyError:  # vacuum
                pass
            else:  # implicit solvent
                tleap.new_section('Set GB radii to recommended values for OBC')
                tleap.add_commands('set default PBRadii {}'.format(
                    get_leap_recommended_pbradii(implicit_solvent)))
        else:  # explicit solvent
            tleap.new_section('Solvate systems')

            # Solvate unit. Solvent models different than tip3p need parameter modifications.
            solvent_model = solvent['solvent_model']
            # Check that solvent model has loaded the appropriate leap parameters.
            # This does guarantee a failure, but is a good sign of it.
            if _OPENMM_LEAP_SOLVENT_FILES_MAP[solvent_model] not in loaded_water_files:
                solvent_warning = ("WARNING: The solvent_model {} may not work for loaded "
                                   "leaprc.water.X files.\n We expected {} to make your "
                                   "solvent model work, but did not find it.\n "
                                   "This is okay for tip4pew leaprc file with tip3p solvent_model, "
                                   "but not the other way around.\nThis does "
                                   "not mean the system will not build, but it may throw "
                                   "an error.".format(solvent_model,
                                                      _OPENMM_LEAP_SOLVENT_FILES_MAP[solvent_model]))
                logger.warning(solvent_warning)
            leap_solvent_model = _OPENMM_LEAP_SOLVENT_MODELS_MAP[solvent_model]
            clearance = solvent['clearance']
            tleap.solvate(unit_name=unit_to_solvate, solvent_model=leap_solvent_model, clearance=clearance)

            # First, determine how many ions we need to add for the ionic strength.
            if not ignore_ionic_strength and solvent['ionic_strength'] != 0.0*unit.molar:
                # Currently we support only monovalent ions.
                for ion_name in [solvent['positive_ion'], solvent['negative_ion']]:
                    assert '2' not in ion_name and '3' not in ion_name

                logger.debug('Estimating number of water molecules in the box.')
                n_waters = self._get_number_box_waters(pack, alchemical_charge, system_parameters,
                                                       solvent_id, *molecule_ids)
                logger.debug('Estimated number of water molecules: {}'.format(n_waters))

                # Water molarity at room temperature: 998.23g/L / 18.01528g/mol ~= 55.41M
                n_ions_ionic_strength = int(np.round(n_waters * solvent['ionic_strength'] / (55.41*unit.molar)))

                logging.debug('Adding {} ions in {} water molecules to reach ionic strength '
                              'of {}'.format(n_ions_ionic_strength, n_waters, solvent['ionic_strength']))
            else:
                n_ions_ionic_strength = 0

            # Add alchemically modified ions that we don't already add for ionic strength.
            if abs(alchemical_charge) > n_ions_ionic_strength:
                n_alchemical_ions = abs(alchemical_charge) - n_ions_ionic_strength
                try:
                    if alchemical_charge > 0:
                        ion = solvent['negative_ion']
                    else:
                        ion = solvent['positive_ion']
                except KeyError:
                    err_msg = ('Found charged ligand but no indications for ions in '
                               'solvent {}').format(solvent_id)
                    logger.error(err_msg)
                    raise RuntimeError(err_msg)
                tleap.add_ions(unit_name=unit_to_solvate, ion=ion,
                               num_ions=n_alchemical_ions, replace_solvent=True)
                logging.debug('Adding {} {} ion to neutralize ligand charge of {}'
                              ''.format(n_alchemical_ions, ion, alchemical_charge))

            # Neutralizing solvation box
            if 'positive_ion' in solvent:
                tleap.add_ions(unit_name=unit_to_solvate, ion=solvent['positive_ion'],
                               replace_solvent=True)
            if 'negative_ion' in solvent:
                tleap.add_ions(unit_name=unit_to_solvate, ion=solvent['negative_ion'],
                               replace_solvent=True)

            # Ions for the ionic strength must be added AFTER neutralization.
            if n_ions_ionic_strength != 0:
                tleap.add_ions(unit_name=unit_to_solvate, ion=solvent['positive_ion'],
                               num_ions=n_ions_ionic_strength, replace_solvent=True)
                tleap.add_ions(unit_name=unit_to_solvate, ion=solvent['negative_ion'],
                               num_ions=n_ions_ionic_strength, replace_solvent=True)

        # Check charge
        tleap.new_section('Check charge')
        tleap.add_commands('check ' + unit_to_solvate)

        # Save output files
        # ------------------
        system_dir = os.path.dirname(system_file_path)
        base_file_path = os.path.basename(system_file_path).split('.')[0]
        base_file_path = os.path.join(system_dir, base_file_path)

        # Create output directory
        if not os.path.exists(system_dir):
            os.makedirs(system_dir)

        # Save prmtop, inpcrd and reference pdb files
        tleap.new_section('Save prmtop and inpcrd files')
        if save_amber_files:
            tleap.save_unit(unit_to_solvate, system_file_path)
        tleap.save_unit(unit_to_solvate, base_file_path + '.pdb')

        # Save tleap script for reference
        tleap.export_script(base_file_path + '.leap.in')

        # Run tleap and log warnings
        # Handle common errors we know of
        try:
            warnings = tleap.run()
        except RuntimeError as e:
            error = RuntimeError('Solvent {}: {}'.format(solvent_id, str(e)))
            error.with_traceback(sys.exc_info()[2])
            raise error

        for warning in warnings:
            logger.warning('TLeap: ' + warning)

    def _get_number_box_waters(self, *args):
        """Build a system in a temporary directory and count the number of waters."""
        with mmtools.utils.temporary_directory() as tmp_dir:
            system_file_path = os.path.join(tmp_dir, 'temp_system.prmtop')
            self._setup_system(system_file_path, *args, ignore_ionic_strength=True,
                               save_amber_files=False)

            # Count number of waters of created system.
            system_file_path = os.path.join(tmp_dir, 'temp_system.pdb')
            system_traj = mdtraj.load(system_file_path)
            n_waters = sum([1 for res in system_traj.topology.residues if res.is_water])

        return n_waters


# ==============================================================================
# FUNCTIONS FOR ALCHEMICAL PATH OPTIMIZATION
# ==============================================================================

class _DCDTrajectoryFile(mdtraj.formats.dcd.DCDTrajectoryFile):
    """Convenience class extending MDTraj DCD trajectory file.

    This handles units and allow reading/writing SamplerStates instead
    of positions, cell_lengths, and cell_angles.

    """

    def read(self, *args, **kwargs):
        positions, cell_lengths, cell_angles = super().read(*args, **kwargs)
        # Add standard DCD units.
        return positions*unit.angstrom, cell_lengths*unit.angstrom, cell_angles*unit.degree

    def read_as_sampler_states(self, *args, **kwargs):
        positions, cell_lengths, cell_angles = self.read(*args, **kwargs)
        sampler_states = []
        for i in range(len(positions)):
            box_vectors = mdtraj.utils.lengths_and_angles_to_box_vectors(
                *(cell_lengths[i]/unit.nanometer),
                *(cell_angles[i]/unit.degree)
            ) * unit.nanometer
            sampler_state = mmtools.states.SamplerState(positions[i], box_vectors=box_vectors)
            sampler_states.append(sampler_state)
        return sampler_states

    def write(self, xyz, cell_lengths=None, cell_angles=None, **kwargs):
        # Convert to standard DCD units.
        super().write(xyz/unit.angstrom, cell_lengths/unit.angstrom, cell_angles/unit.degree)

    def write_sampler_state(self, sampler_state):
        a, b, c, alpha, beta, gamma = mdtraj.utils.box_vectors_to_lengths_and_angles(
            *(np.array(sampler_state.box_vectors / unit.angstrom)))
        super().write(sampler_state.positions / unit.angstrom,
                      (a, b, c), (alpha, beta, gamma))


def read_trailblaze_checkpoint_coordinates(checkpoint_dir_path, redistributed=True):
    """Read positions and box vectors stored as checkpoint by the trailblaze algorithm.

    Parameters
    ----------
    checkpoint_dir_path : str
        The path to the directory containing the checkpoint information.
    redistributed : bool, optional
        If True, the function will check if the states were redistributed,
        and will returned the set of coordinates that are more representative
        of the redistributed protocol.

    Returns
    -------
    sampler_states : List[openmmtools.states.SamplerState], optional
        ``sampler_states[i]`` contain positions and box vectors for the
        intermediate state i generated by the trailblaze algorithm.

    Raises
    ------
    FileNotFoundError
        If no file with the coordinates was found.
    """
    positions_file_path = os.path.join(checkpoint_dir_path, 'coordinates.dcd')
    states_map_file_path = os.path.join(checkpoint_dir_path, 'states_map.json')

    # Open the file if it exist.
    try:
        trajectory_file = _DCDTrajectoryFile(positions_file_path, 'r')
    except OSError as e:
        raise FileNotFoundError(str(e))

    # Read info.
    try:
        sampler_states = trajectory_file.read_as_sampler_states()
    finally:
        trajectory_file.close()

    # If the protocol was redistributed, use the states map to create
    # a new set of sampler states that can be used as starting conditions
    # for the redistributed protocol.
    if redistributed:
        try:
            with open(states_map_file_path, 'r') as f:
                states_map = json.load(f)
            sampler_states = [sampler_states[i] for i in states_map]
        except FileNotFoundError:
            pass

    return sampler_states


def _resume_thermodynamic_trailblazing(checkpoint_dir_path, initial_protocol):
    """Resume a previously-run trailblaze execution.

    Parameters
    ----------
    checkpoint_dir_path : str
        The path to the directory used to store the trailblaze information.
    initial_protocol : Dict[str, List[float]]
        The initial protocol containing only the first state of the path.
        If no checkpoint file for the protocol is found, the file will
        be initialized using this state.

    Returns
    -------
    resumed_protocol : Dict[str, List[float]]
        The resumed optimal protocol.
    trajectory_file : _DCDTrajectoryFile
        A DCD file open for appening.
    sampler_state : SamplerState or None
        The last saved SamplerState or None if no frame was saved yet.

    """
    # We save the protocol in a YAML file and the positions in netcdf.
    protocol_file_path = os.path.join(checkpoint_dir_path, 'protocol.yaml')
    stds_file_path = os.path.join(checkpoint_dir_path, 'states_stds.json')
    positions_file_path = os.path.join(checkpoint_dir_path, 'coordinates.dcd')

    # Create the directory, if it doesn't exist.
    os.makedirs(checkpoint_dir_path, exist_ok=True)

    # Load protocol and stds checkpoint file.
    try:
        # Parse the previously calculated optimal_protocol dict.
        with open(protocol_file_path, 'r') as file_stream:
            resumed_protocol = yaml.load(file_stream, Loader=yaml.FullLoader)

        # Load the energy difference stds.
        with open(stds_file_path, 'r') as f:
            states_stds = json.load(f)
    except FileNotFoundError:
        resumed_protocol = initial_protocol
        states_stds = [[], []]

    # Check if there's an existing positions information.
    try:
        # We want the coordinates of the states that were sampled
        # during the search not the states after redistribution.
        sampler_states = read_trailblaze_checkpoint_coordinates(
            checkpoint_dir_path, redistributed=False)
    except FileNotFoundError:
        len_trajectory = 0
    else:
        len_trajectory = len(sampler_states)

    # Raise an error if the algorithm was interrupted *during*
    # writing on disk. We store only the states that we simulated
    # so there should be one less here.
    for state_values in resumed_protocol.values():
        if len_trajectory < len(state_values) - 1:
            err_msg = ("The trailblaze algorithm was interrupted while "
                       "writing the checkpoint file and it is now unable "
                       "to resume. Please delete the files "
                       f"in {checkpoint_dir_path} and restart.")
            raise RuntimeError(err_msg)

    # When this is resumed, but the trajectory is already completed,
    # the frame of the final end state has been already written, but
    # we don't want to add it twice at the end of the trailblaze function.
    len_trajectory = len(state_values) - 1

    # Whether the file exist or not, MDTraj doesn't support appending
    # files so we open a new one and rewrite the configurations we
    # have generated in the previous run.
    trajectory_file = _DCDTrajectoryFile(positions_file_path, 'w',
                                         force_overwrite=True)
    if len_trajectory > 0:
        for i in range(len_trajectory):
            trajectory_file.write_sampler_state(sampler_states[i])
        # Make sure the next search starts from the last saved position
        # unless the previous calculation was interrupted before the
        # first position could be saved.
        sampler_state = sampler_states[-1]
    else:
        sampler_state = None

    return resumed_protocol, states_stds, trajectory_file, sampler_state


def _cache_trailblaze_data(checkpoint_dir_path, optimal_protocol, states_stds,
                           trajectory_file, sampler_state):
    """Store on disk current state of the trailblaze run."""

    # Determine the file paths of the stored data.
    protocol_file_path = os.path.join(checkpoint_dir_path, 'protocol.yaml')
    stds_file_path = os.path.join(checkpoint_dir_path, 'states_stds.json')

    # Update protocol.
    with open(protocol_file_path, 'w') as file_stream:
        yaml.dump(optimal_protocol, file_stream)

    # Update the stds between states.
    with open(stds_file_path, 'w') as f:
        json.dump(states_stds, f)

    # Append positions of the state that we just simulated.
    trajectory_file.write_sampler_state(sampler_state)


def _redistribute_trailblaze_states(old_protocol, states_stds, thermodynamic_distance):
    """Redistribute the states using a bidirectional estimate of the thermodynamic length.

    Parameters
    ----------
    old_protocol : Dict[str, List[float]]
        The unidirectional optimal protocol.
    states_stds : List[List[float]]
        states_stds[j][i] is the standard deviation of the potential
        difference between states i-1 and i computed in the j direction.
    thermodynamic_distance : float
        The distance between each pair of states.

    Returns
    -------
    new_protocol : Dict[str, List[float]]
        The new estimate of the optimal protocol.
    states_map : List[int]
        states_map[i] is the index of the state in the old protocol that
        is closest to the i-th state in the new protocol. This allows to
        map coordinates generated during trailblazing to the redistributed
        protocol.
    """
    # The parameter names in a fixed order.
    parameter_names = [par_name for par_name in old_protocol]

    # Initialize the new protocol from the first state the optimal protocol.
    new_protocol = {par_name: [values[0]] for par_name, values in old_protocol.items()}

    # The first state of the new protocol always maps to the first state of the old one.
    states_map = [0]

    def _get_old_protocol_state(state_idx):
        """Return a representation of the thermo state as a list of parameter values."""
        return np.array([old_protocol[par_name][state_idx] for par_name in parameter_names])

    def _add_state_to_new_protocol(state):
        for parameter_name, new_state_value in zip(parameter_names, state.tolist()):
            new_protocol[parameter_name].append(new_state_value)

    # The thermodynamic length at 0 is 0.0.
    states_stds[0] = [0.0] + states_stds[0]
    states_stds[1] = [0.0] + states_stds[1]

    # We don't have the energy difference std in the
    # direction opposite to the search direction so we
    # pad the list.
    states_stds[1].append(states_stds[0][-1])
    states_stds = np.array(states_stds)

    # Compute a bidirectional estimate of the thermodynamic length.
    old_protocol_thermo_length = np.cumsum(np.mean(states_stds, axis=0))

    # Trailblaze again interpolating the thermodynamic length function.
    current_state_idx = 0
    current_state = _get_old_protocol_state(0)
    last_state = _get_old_protocol_state(-1)
    new_protocol_cum_thermo_length = 0.0
    while (current_state != last_state).any():
        # Find first state for which the accumulated standard
        # deviation is greater than the thermo length threshold.
        try:
            while old_protocol_thermo_length[current_state_idx+1] - new_protocol_cum_thermo_length <= thermodynamic_distance:
                current_state_idx += 1
        except IndexError:
            # If we got to the end, we just add the last state
            # to the protocol and stop the while loop.
            _add_state_to_new_protocol(last_state)
            break

        # Update current state.
        current_state = _get_old_protocol_state(current_state_idx)

        # The thermodynamic length from the last redistributed state to current state.
        pair_thermo_length = old_protocol_thermo_length[current_state_idx] - new_protocol_cum_thermo_length

        # Now interpolate between the current state state and the next to find
        # the exact state for which the thermo length equal the threshold.
        next_state = _get_old_protocol_state(current_state_idx+1)
        differential = thermodynamic_distance - pair_thermo_length
        differential /= old_protocol_thermo_length[current_state_idx+1] - old_protocol_thermo_length[current_state_idx]
        new_state = current_state + differential * (next_state - current_state)

        # Update cumulative thermo length.
        new_protocol_cum_thermo_length += thermodynamic_distance

        # Update states map.
        closest_state_idx = current_state_idx if differential <= 0.5 else current_state_idx+1
        states_map.append(closest_state_idx)

        # Update redistributed protocol.
        _add_state_to_new_protocol(new_state)

    # The last state of the new protocol always maps to the last state of the old one.
    states_map.append(len(old_protocol_thermo_length)-1)

    return new_protocol, states_map


def run_thermodynamic_trailblazing(
        thermodynamic_state, sampler_state, mcmc_move, state_parameters,
        parameter_setters=None, thermodynamic_distance=1.0,
        distance_tolerance=0.05, n_samples_per_state=100,
        reversed_direction=False, bidirectional_redistribution=True,
        bidirectional_search_thermo_dist='auto',
        global_parameter_functions=None, function_variables=tuple(),
        checkpoint_dir_path=None
):
    """
    Find an alchemical path by placing alchemical states at a fixed distance.

    The distance between two states is estimated by collecting ``n_samples_per_state``
    configurations through the MCMCMove in one of the two alchemical states,
    and computing the standard deviation of the difference of potential energies
    between the two states at those configurations.

    The states of the protocol are chosen so that each pair has a distance
    (in thermodynamic length) of ``thermodynamic_distance +- distance_tolerance``.
    The thermodynamic length estimate (in kT) is based on the standard deviation
    of the difference in potential energy between the two states.

    The function is capable of resuming when interrupted if ``checkpoint_dir_path``
    is specified. This will create two files called 'protocol.yaml' and
    'coordinates.dcd' storing the protocol and initial positions and box
    vectors for each state that are generated while running the algorithm.

    It is also possible to discretize a path specified through maathematical
    expressions through the arguments ``global_parameter_function`` and
    ``function_variables``.

    Parameters
    ----------
    thermodynamic_state : openmmtools.states.CompoundThermodynamicState
        The state of the alchemically modified system.
    sampler_state : openmmtools.states.SamplerState
        The sampler states including initial positions and box vectors.
    mcmc_move : openmmtools.mcmc.MCMCMove
        The MCMCMove to use for propagation.
    state_parameters : List[Tuple[str, List[float]]]
        Each element of this list is a tuple containing first the name
        of the parameter to be modified (e.g. ``lambda_electrostatics``,
        ``lambda_sterics``) and a list specifying the initial and final
        values for the path.
    parameter_setters : Dict[str, Callable], optional
        If the parameter cannot be set in the ``thermodynamic_state``
        with a simple call to ``setattr``, you can pass a dictionary
        mapping the parameter name to a function
        ``setter(thermodynamic_state, parameter_name, value)``. This
        is useful for example to set global parameter function variables
        with ``openmmtools.states.GlobalParameterState.set_function_variable``.
    thermodynamic_distance : float, optional
        The target distance (in thermodynamic length) between each pair of
        states in kT. Default is 1.0 (kT).
    distance_tolerance : float, optional
        The tolerance on the found standard deviation. Default is 0.05 (kT).
    n_samples_per_state : int, optional
        How many samples to collect to estimate the overlap between two
        states. Default is 100.
    reversed_direction : bool, optional
        If ``True``, the algorithm starts from the final state and traverses
        the path from the end to the beginning. The returned path
        discretization will still be ordered from the beginning to the
        end following the order in ``state_parameters``. Default is ``False``.
    bidirectional_redistribution : bool, optional
        If ``True``, the states will be redistributed using the standard
        deviation of the potential difference between states in both
        directions. Default is ``True``.
    bidirectional_search_thermo_dist : float or 'auto', optional
        If ``bidirectional_redistribution`` is ``True``, the thermodynamic
        distance between the sampled states used to collect data along
        the path can be different than the thermodynamic distance after
        redistribution. The default ('auto') caps the thermodynamic
        distance used for trailblazing at 1 kT. Keeping this value small
        lower the chance of obtaining very large stds in the opposite direction
        due to rare, dominating events in sections of the path where the overlap
        decreases quickly, which in turn may results in unreasonably long
        protocols.
    global_parameter_functions : Dict[str, Union[str, openmmtools.states.GlobalParameterFunction]], optional
        Map a parameter name to a mathematical expression as a string
        or a ``openmmtools.states.GlobalParameterFunction`` object.
    function_variables : List[str], optional
        A list of function variables entering the mathematical
        expressions.
    checkpoint_dir_path : str, optional
        The path to the directory used to store the trailblaze information.
        If this is given and the files exist, the algorithm will use this
        information to resume in case it was previously interrupted. If
        ``None``, no information is stored and it won't be possible to
        resume. Default is ``None``.

    Returns
    -------
    optimal_protocol : Dict[str, List[float]]
        The estimated protocol. Each dictionary key is one of the
        parameters in ``state_parameters``, and its values is the
        list of values that it takes in each state of the path.

    """
    def _state_parameter_setter(state, parameter_name, value):
        """Helper function to set state parameters."""
        setattr(state, parameter_name, value)

    def _function_variable_setter(state, parameter_name, value):
        """Helper function to set global parameter function variables."""
        state.set_function_variable(parameter_name, value)

    # Make sure that the state parameters to optimize have a clear order.
    assert (isinstance(state_parameters, list) or isinstance(state_parameters, tuple))

    # Determine the thermo distance to achieve during the search.
    if not bidirectional_redistribution:
        search_thermo_dist = thermodynamic_distance
    else:
        if bidirectional_search_thermo_dist == 'auto':
            search_thermo_dist = min(1.0, thermodynamic_distance)
        else:
            search_thermo_dist = bidirectional_search_thermo_dist

    # Create unordered helper variable.
    state_parameter_dict = {x[0]: x[1] for x in state_parameters}

    # Do not modify original thermodynamic_state.
    thermodynamic_state = copy.deepcopy(thermodynamic_state)

    # Handle mutable default arguments.
    if parameter_setters is None:
        parameter_setters = {}
    if global_parameter_functions is None:
        global_parameter_functions = {}

    # Make sure that the same parameter was not listed both as
    # a function and a parameter to iterate.
    for parameter_name, _ in state_parameters:
        if parameter_name in global_parameter_functions:
            raise ValueError(f"Cannot specify {parameter_name} in "
                              "'state_parameters' and 'global_parameter_functions'")

    # Make sure all function variables are in state_parameters.
    for function_variable in function_variables:
        if function_variable not in state_parameter_dict:
            raise ValueError(f"The variable '{function_variable}' must be given in 'state_parameters'")

    # Use special setters for the function variables.
    for function_variable in function_variables:
        if parameter_name not in parameter_setters:
            parameter_setters[function_variable] = _function_variable_setter
    # Create default setters for all other state parameters to later avoid if/else or try/except.
    for parameter_name, _ in state_parameters:
        if parameter_name not in parameter_setters:
            parameter_setters[parameter_name] = _state_parameter_setter

    # Initialize all function variables in the thermodynamic state.
    for function_variable in function_variables:
        value = state_parameter_dict[function_variable][0]
        thermodynamic_state.set_function_variable(function_variable, value)

    # Assign global parameter functions.
    for parameter_name, global_parameter_function in global_parameter_functions.items():
        # If the user doesn't pass an instance of function class, create one.
        if isinstance(global_parameter_function, str):
            global_parameter_function = mmtools.states.GlobalParameterFunction(global_parameter_function)
        setattr(thermodynamic_state, parameter_name, global_parameter_function)

    # Reverse the direction of the algorithm if requested.
    if reversed_direction:
        state_parameters = [(par_name, end_states.__class__(reversed(end_states)))
                            for par_name, end_states in reversed(state_parameters)]

    # Initialize protocol with the starting value.
    optimal_protocol = {par: [values[0]] for par, values in state_parameters}

    # Keep track of potential std between states in both directions
    # of the path so that we can redistribute the states later.
    # At the end of the protocol this will have the same length
    # of the protocol minus one. The inner lists are for the forward
    # and reversed direction stds respectively.
    states_stds = [[], []]

    # Check to see whether a trailblazing algorithm is already in progress,
    # and if so, restore to the previously checkpointed state.
    if checkpoint_dir_path is not None:
        optimal_protocol, states_stds, trajectory_file, resumed_sampler_state = _resume_thermodynamic_trailblazing(
            checkpoint_dir_path, optimal_protocol)
        # Start from the last saved conformation.
        if resumed_sampler_state is not None:
            sampler_state = resumed_sampler_state

    # We keep track of the previous state in the optimal protocol
    # that we'll use to compute the stds in the opposite direction.
    if len(states_stds[0]) == 0:
        previous_thermo_state = None
    else:
        previous_thermo_state = copy.deepcopy(thermodynamic_state)

    # Make sure that thermodynamic_state is in the last explored
    # state, whether the algorithm was resumed or not.
    for state_parameter in optimal_protocol:
        parameter_setters[state_parameter](thermodynamic_state, state_parameter,
                                           optimal_protocol[state_parameter][-1])
        if previous_thermo_state is not None:
            parameter_setters[state_parameter](previous_thermo_state, state_parameter,
                                               optimal_protocol[state_parameter][-2])

    # We change only one parameter at a time.
    for state_parameter, values in state_parameters:
        logger.debug('Determining alchemical path for parameter {}'.format(state_parameter))

        # Is this a search from 0 to 1 or from 1 to 0?
        search_direction = np.sign(values[1] - values[0])
        # If the parameter doesn't change, continue to the next one.
        if search_direction == 0:
            continue

        # Gather data until we get to the last value.
        while optimal_protocol[state_parameter][-1] != values[-1]:
            # Simulate current thermodynamic state to collect samples.
            sampler_states = []
            simulated_energies = np.zeros(n_samples_per_state)
            for i in range(n_samples_per_state):
                mcmc_move.apply(thermodynamic_state, sampler_state)
                sampler_states.append(copy.deepcopy(sampler_state))

            # Keep track of the thermo state we use for the reweighting.
            reweighted_thermo_state = None

            # Find first state that doesn't overlap with simulated one
            # with std(du) within search_thermo_dist +- distance_tolerance.
            # We stop anyway if we reach the last value of the protocol.
            std_energy = 0.0
            current_parameter_value = optimal_protocol[state_parameter][-1]
            while (abs(std_energy - search_thermo_dist) > distance_tolerance and
                   not (current_parameter_value == values[1] and std_energy < search_thermo_dist)):

                # Determine next parameter value to compute.
                if np.isclose(std_energy, 0.0):
                    # This is the first iteration or the two state overlap significantly
                    # (e.g. small molecule in vacuum). Just advance by a +- 0.05 step.
                    old_parameter_value = current_parameter_value
                    current_parameter_value += (values[1] - values[0]) / 20.0
                else:
                    # Assume std_energy(parameter_value) is linear to determine next value to try.
                    derivative_std_energy = ((std_energy - old_std_energy) /
                                             (current_parameter_value - old_parameter_value))
                    old_parameter_value = current_parameter_value
                    current_parameter_value += (search_thermo_dist - std_energy) / derivative_std_energy

                # Keep current_parameter_value inside bound interval.
                if search_direction * current_parameter_value > values[1]:
                    current_parameter_value = values[1]
                assert search_direction * (optimal_protocol[state_parameter][-1] - current_parameter_value) < 0

                # Determine the thermo states at which we need to compute the energies.
                # If this is the first attempt, compute also the reduced potential of
                # the simulated energies and the previous state to estimate the standard
                # deviation in the opposite direction.
                if reweighted_thermo_state is None:
                    # First attempt.
                    reweighted_thermo_state = copy.deepcopy(thermodynamic_state)
                    computed_thermo_states = [reweighted_thermo_state, thermodynamic_state]
                    if previous_thermo_state is not None:
                        computed_thermo_states.append(previous_thermo_state)
                else:
                    computed_thermo_states = [reweighted_thermo_state]

                # Set the reweighted state to the current parameter value.
                parameter_setters[state_parameter](reweighted_thermo_state, state_parameter, current_parameter_value)

                # Compute all energies.
                energies = np.empty(shape=(len(computed_thermo_states), n_samples_per_state))
                for i, sampler_state in enumerate(sampler_states):
                    energies[:,i] = mmtools.states.reduced_potential_at_states(
                        sampler_state, computed_thermo_states, mmtools.cache.global_context_cache)

                # Cache the simulated energies for the next iteration.
                if len(computed_thermo_states) > 1:
                    simulated_energies = energies[1]

                # Compute the energy difference std in the direction: simulated state -> previous state.
                if len(computed_thermo_states) > 2:
                    denergies = energies[2] - simulated_energies
                    states_stds[1].append(float(np.std(denergies, ddof=1)))

                # Compute the energy difference std between the currently simulated and the reweighted states.
                old_std_energy = std_energy
                denergies = energies[0] - simulated_energies
                std_energy = np.std(denergies, ddof=1)
                logger.debug('trailblazing: state_parameter {}, simulated_value {}, current_parameter_value {}, '
                             'std_du {}'.format(state_parameter, optimal_protocol[state_parameter][-1],
                                                current_parameter_value, std_energy))

            # Store energy difference std in the direction: simulated state -> reweighted state.
            states_stds[0].append(float(std_energy))

            # Update variables for next iteration.
            previous_thermo_state = copy.deepcopy(thermodynamic_state)
            thermodynamic_state = reweighted_thermo_state

            # Update the optimal protocol with the new value of this parameter.
            # The other parameters remain fixed.
            for par_name in optimal_protocol:
                # Make sure we append to a Python float to the list.
                # Lists of numpy types sometimes give problems.
                if par_name == state_parameter:
                    protocol_value = float(current_parameter_value)
                else:
                    protocol_value = float(optimal_protocol[par_name][-1])
                optimal_protocol[par_name].append(protocol_value)

            # Save the updated checkpoint file to disk.
            if checkpoint_dir_path is not None:
                _cache_trailblaze_data(checkpoint_dir_path, optimal_protocol, states_stds,
                                       trajectory_file, sampler_state)

    if checkpoint_dir_path is not None:
        # We haven't simulated the last state so we just set the positions of the second to last.
        trajectory_file.write_sampler_state(sampler_state)
        trajectory_file.close()

    # Redistribute the states using the standard deviation estimates in both directions.
    if bidirectional_redistribution:
        optimal_protocol, states_map = _redistribute_trailblaze_states(
            optimal_protocol, states_stds, thermodynamic_distance)

    # Save the states map for reading the coordinates correctly.
    if checkpoint_dir_path is not None:
        states_map_file_path = os.path.join(checkpoint_dir_path, 'states_map.json')

        if bidirectional_redistribution:
            with open(states_map_file_path, 'w') as f:
                json.dump(states_map, f)
        elif os.path.isfile(states_map_file_path):
            # If there's an old file because the path was previously redistributed,
            # we delete it so that read_trailblaze_checkpoint_coordinates will
            # return the coordinates associated to the most recently-generated path.
            os.remove(states_map_file_path)

    # If we have traversed the path in the reversed direction, re-invert
    # the order of the discretized path.
    if reversed_direction:
        for par_name, values in optimal_protocol.items():
            optimal_protocol[par_name] = values.__class__(reversed(values))

    # If we used global parameter functions, the optimal_protocol at this
    # point is a discretization of the function_variables, not the original
    # parameters so we convert it back.
    len_protocol = len(next(iter(optimal_protocol.values())))
    function_variables_protocol = {var: optimal_protocol.pop(var) for var in function_variables}
    original_parameters_protocol = {par: [] for par in global_parameter_functions}

    # Rebuild the optimal discretization for the original parameters.
    for state_idx in range(len_protocol):
        # Set the function variable value so that the function is computed.
        for function_variable in function_variables:
            value = function_variables_protocol[function_variable][state_idx]
            thermodynamic_state.set_function_variable(function_variable, value)

        # Recover original parameters.
        for parameter_name in original_parameters_protocol:
            value = getattr(thermodynamic_state, parameter_name)
            original_parameters_protocol[parameter_name].append(value)

    # Update the total protocol.
    optimal_protocol.update(original_parameters_protocol)

    logger.debug('Alchemical path found: {}'.format(optimal_protocol))
    return optimal_protocol


if __name__ == '__main__':
    import doctest
    doctest.testmod()
