"""
This file contains the base classes for topology proposals
"""

from simtk import unit

from simtk import openmm
from simtk.openmm import app

from collections import namedtuple
import copy
import warnings
import logging
import itertools
import json
import os
import openeye.oechem as oechem
import numpy as np
import openeye.oeomega as oeomega
import tempfile
import networkx as nx
from openmoltools import forcefield_generators
import openeye.oegraphsim as oegraphsim
from perses.rjmc.geometry import FFAllAngleGeometryEngine
from perses.storage import NetCDFStorageView
from io import StringIO
import openmoltools
import base64
import progressbar
from typing import List, Dict
from perses.rjmc.geometry import NoTorsionError
try:
    from subprocess import getoutput  # If python 3
except ImportError:
    from commands import getoutput  # If python 2

################################################################################
# CONSTANTS
################################################################################

OESMILES_OPTIONS = oechem.OESMILESFlag_DEFAULT | oechem.OESMILESFlag_ISOMERIC | oechem.OESMILESFlag_Hydrogens

# TODO write a mapping-protocol class to handle these options

# weak requirements for mapping atoms == more atoms mapped, more in core
# atoms need to match in aromaticity. Same with bonds.
# maps ethane to ethene, CH3 to NH2, but not benzene to cyclohexane
WEAK_ATOM_EXPRESSION = oechem.OEExprOpts_EqAromatic | oechem.OEExprOpts_EqNotAromatic | oechem.OEExprOpts_IntType
WEAK_BOND_EXPRESSION = oechem.OEExprOpts_DefaultBonds

# default atom expression, requires same aromaticitiy and hybridization
# bonds need to match in bond order
# ethane to ethene wouldn't map, CH3 to NH2 would map but CH3 to HC=O wouldn't
DEFAULT_ATOM_EXPRESSION = oechem.OEExprOpts_Hybridization | oechem.OEExprOpts_IntType
DEFAULT_BOND_EXPRESSION = oechem.OEExprOpts_DefaultBonds

# strong requires same hybridization AND the same atom type
# bonds are same as default, require them to match in bond order
STRONG_ATOM_EXPRESSION = oechem.OEExprOpts_Hybridization | oechem.OEExprOpts_HvyDegree | oechem.OEExprOpts_DefaultAtoms
STRONG_BOND_EXPRESSION = oechem.OEExprOpts_DefaultBonds

################################################################################
# LOGGER
################################################################################

import logging
logging.basicConfig(level = logging.NOTSET)
_logger = logging.getLogger("proposal_generator")
_logger.setLevel(logging.INFO)

################################################################################
# UTILITIES
################################################################################

def append_topology(destination_topology, source_topology, exclude_residue_name=None):
    """
    Add the source OpenMM Topology to the destination Topology.

    Parameters
    ----------
    destination_topology : simtk.openmm.app.Topology
        The Topology to which the contents of `source_topology` are to be added.
    source_topology : simtk.openmm.app.Topology
        The Topology to be added.
    exclude_residue_name : str, optional, default=None
        If specified, any residues matching this name are excluded.

    """
    if exclude_residue_name is None:
        exclude_residue_name = "   " #something with 3 characters that is never a residue name
    new_atoms = {}
    for chain in source_topology.chains():
        new_chain = destination_topology.addChain(chain.id)
        for residue in chain.residues():
            # TODO: should we use complete residue names?
            if (residue.name[:3] == exclude_residue_name[:3]):
                continue
            new_residue = destination_topology.addResidue(residue.name, new_chain, residue.id)
            for atom in residue.atoms():
                new_atom = destination_topology.addAtom(atom.name, atom.element, new_residue, atom.id)
                new_atoms[atom] = new_atom
    for bond in source_topology.bonds():
        if (bond[0].residue.name[:3] == exclude_residue_name[:3]) or (bond[1].residue.name[:3] == exclude_residue_name[:3]):
            continue
        order = bond.order
        destination_topology.addBond(new_atoms[bond[0]], new_atoms[bond[1]], order=order)

def deepcopy_topology(source_topology):
    """
    Drop-in replacement for copy.deepcopy(topology) that fixes backpointer issues.

    Parameters
    ----------
    source_topology : simtk.openmm.app.Topology
        The Topology to be added.

    """
    topology = app.Topology()
    append_topology(topology, source_topology)
    return topology

def has_h_mapped(atommap, mola: oechem.OEMol, molb: oechem.OEMol):
    for a_atom, b_atom in atommap.items():
        if mola.GetAtom(oechem.OEHasAtomIdx(a_atom)).GetAtomicNum() == 1 or molb.GetAtom(oechem.OEHasAtomIdx(b_atom)).GetAtomicNum() == 1:
            return True

    return False

class InvalidMappingException(Exception):
    def __init__(self, message, errors):

        # Call the base class constructor with the parameters it needs
        super().__init__(message)

        # Now for your custom code...
        self.errors = errors

class AtomMapper(object):
    """ Generates map between two molecules.
    As this doesn't generate a system, it will not be
    accurate for whether hydrogens are mapped or not.
    It also doesn't check that this is feasible to simulate,
    but is just helpful for testing options.

    Parameters
    ----------
    current_molecule : oechem.OEMol
        starting molecule
    proposed_molecule : oechem.OEMol
        molecule to move to
    atom_expr : oechem.OEExprOpts
        atom requirement to consider two bonds as the same
    bond_expr : oechem.OEExprOpts
        bond requirement to consider two bonds as the same
    verbose : bool
        wether to be verbose or not
    allow_ring_breaking : bool, default=True
        Wether or not to allow ring breaking in map

    Attributes
    ----------
    map : type
        Description of attribute `map`.
    _get_mol_atom_map : type
        Description of attribute `_get_mol_atom_map`.
    current_molecule
    proposed_molecule
    atom_expr
    bond_expr
    verbose
    allow_ring_breaking

    """
    def __init__(self, list_of_oemols, atom_expr=None, bond_expr=None,
                 map_strength='default', verbose=False, allow_ring_breaking=False, **kwargs):
        self.list_of_oemols = list_of_oemols
        self.map_strength = map_strength
        self.atom_expr = atom_expr
        self.bond_expr = bond_expr
        self.verbose = verbose
        self.allow_ring_breaking = allow_ring_breaking


        self.atom_map = self._get_mol_atom_map()
        super(AtomMapper, self).__init__(**kwargs)

    def _get_mol_atom_map(self, indexA=0, indexB=1, allow_ring_breaking=False):
        """
        Given two molecules, returns the mapping of atoms between them using the match with the greatest number of atoms

        Arguments
        ---------
        current_molecule : openeye.oechem.oemol object
             The current molecule in the sampler
        proposed_molecule : openeye.oechem.oemol object
             The proposed new molecule
        allow_ring_breaking : bool, optional, default=False
             If False, will check to make sure rings are not being broken or formed.

        Returns
        -------
        matches : list of match
            list of the matches between the molecules, or None if no matches possible

        """
        self.current_molecule = self.list_of_oemols[indexA]
        self.proposed_molecule = self.list_of_oemols[indexB]
        # TODO this needs to handle all string inputs that might be used
        if self.atom_expr is not None:
            _logger.debug(f'Using user defined atom_expr')
            _logger.debug('Any map_strength set has been ignored')
        else:
            _logger.info(f'Map strength of {self.map_strength}')
            # setting atom expr using map_strength
            if self.map_strength == 'default':
                self.atom_expr = DEFAULT_ATOM_EXPRESSION
            elif self.map_strength == 'weak':
                self.atom_expr = WEAK_ATOM_EXPRESSION
            elif self.map_strength == 'strong':
                self.atom_expr = STRONG_ATOM_EXPRESSION
            else:
                _logger.warning(f"User defined map_strength: {map_strength} not recognised, setting to default")
                self.atom_expr = DEFAULT_ATOM_EXPRESSION
        if self.bond_expr is not None:
            _logger.debug(f'Using user defined bond_expr')
            _logger.debug('Any map_strength set has been ignored')
        else:
            # setting atom expr using map_strength
            if self.map_strength == 'default':
                self.bond_expr = DEFAULT_BOND_EXPRESSION
            elif self.map_strength == 'weak':
                self.bond_expr = WEAK_BOND_EXPRESSION
            elif self.map_strength == 'strong':
                self.bond_expr = STRONG_BOND_EXPRESSION
            else:
                #_logger.warning(f"User defined map_strength: {map_strength} not recognised, setting to default")
                self.bond_expr = DEFAULT_BOND_EXPRESSION

        # this ensures that the hybridization of the oemols is done for correct atom mapping
        oechem.OEAssignHybridization(self.current_molecule)
        oechem.OEAssignHybridization(self.proposed_molecule)
        self.oegraphmol_current = oechem.OEGraphMol(self.current_molecule) # pattern molecule
        self.oegraphmol_proposed = oechem.OEGraphMol(self.proposed_molecule) # target molecule
        if not self.allow_ring_breaking:
            # assigning ring membership to prevent ring breaking
            self._assign_ring_ids(self.oegraphmol_current)
            self._assign_ring_ids(self.oegraphmol_proposed)
        mcs = oechem.OEMCSSearch(oechem.OEMCSType_Approximate)
        mcs.Init(self.oegraphmol_current, self.atom_expr, self.bond_expr)
        mcs.SetMCSFunc(oechem.OEMCSMaxBondsCompleteCycles())
        unique = False
        matches = [m for m in mcs.Match(self.oegraphmol_proposed, unique)]

        if self.allow_ring_breaking is False:
            # Filter the matches to remove any that allow ring breaking
            matches = [m for m in matches if AtomMapper.preserves_rings(m, self.current_molecule, self.proposed_molecule)]

        if not matches:
            _logger.info('Cannot generate atom map without breaking rings, trying again with weaker mapping.')
            mcs.SetMCSFunc(oechem.OEMCSMaxAtoms())
            unique = False
            matches = [m for m in mcs.Match(self.oegraphmol_proposed, unique)]

            if self.allow_ring_breaking is False:
                # Filter the matches to remove any that allow ring breaking
                matches = [m for m in matches if AtomMapper.preserves_rings(m, self.current_molecule, self.proposed_molecule)]
            if not matches:
                return None

        top_matches = self.rank_degenerate_maps(matches, self.current_molecule, self.proposed_molecule) #remove the matches with the lower rank score (filter out bad degeneracies)
        _logger.debug(f"\tthere are {len(top_matches)} top matches")
        max_num_atoms = max([match.NumAtoms() for match in top_matches])
        _logger.debug(f"\tthe max number of atom matches is: {max_num_atoms}; there are {len([m for m in top_matches if m.NumAtoms() == max_num_atoms])} matches herein")
        new_top_matches = [m for m in top_matches if m.NumAtoms() == max_num_atoms]
        new_to_old_atom_maps = [AtomMapper.hydrogen_mapping_exceptions(match, self.current_molecule, self.proposed_molecule) for match in new_top_matches]
        _logger.debug(f"\tnew to old atom maps with most atom hits: {new_to_old_atom_maps}")

        #now all else is equal; we will choose the map with the highest overlap of atom indices
        index_overlap_numbers = []
        for map in new_to_old_atom_maps:
            hit_number = 0
            for key, value in map.items():
                if key == value:
                    hit_number += 1
            index_overlap_numbers.append(hit_number)

        max_index_overlap_number = max(index_overlap_numbers)
        max_index = index_overlap_numbers.index(max_index_overlap_number)
        _logger.debug(f"\tchose {new_to_old_atom_maps[max_index]}")

        return new_to_old_atom_maps[max_index]

    @staticmethod
    def hydrogen_mapping_exceptions(match, current, proposed):
        """
        Returns an atom map that omits hydrogen-to-nonhydrogen atom maps AND X-H to Y-H where element(X) != element(Y)
        or aromatic(X) != aromatic(Y)
        """
        new_to_old_atom_map = {}

        for matchpair in match.GetAtoms():
            old_index, new_index = matchpair.pattern.GetIdx(), matchpair.target.GetIdx()
            old_atom, new_atom = current.GetAtom(oechem.OEHasAtomIdx(old_index)), proposed.GetAtom(oechem.OEHasAtomIdx(new_index))

            #Check if a hydrogen was mapped to a non-hydroden (basically the xor of is_h_a and is_h_b)
            if (old_atom.GetAtomicNum() == 1) != (new_atom.GetAtomicNum() == 1):
                continue

            new_to_old_atom_map[new_index] = old_index
        return new_to_old_atom_map

    def _assign_ring_ids(self, molecule, max_ring_size=10):
        """ Sets the Int of each atom in the oemol to a number
        corresponding to the ring membership of that atom

        Parameters
        ----------
        molecule : oechem.OEMol
            oemol to assign ring ID to
        max_ring_size : int, default = 10
            Largest ring size that will be checked for

        Returns
        -------
        """
        for atom in molecule.GetAtoms():
            rings = ''
            for i in range(3, max_ring_size+1): # smallest feasible ring size is 3
                rings += str(int(oechem.OEAtomIsInRingSize(atom, i)))
            ring_as_base_ten = int(rings, 2)
            atom.SetIntType(ring_as_base_ten)

    @staticmethod
    def preserves_rings(match, current, proposed):
        """Returns True if the transformation allows ring
        systems to be broken or created."""
        pattern_atoms = {atom.GetIdx(): atom for atom in current.GetAtoms()}
        target_atoms = {atom.GetIdx(): atom for atom in proposed.GetAtoms()}
        pattern_to_target_map = {pattern_atoms[matchpair.pattern.GetIdx()]:
                                 target_atoms[matchpair.target.GetIdx()]
                                 for matchpair in match.GetAtoms()}
        if AtomMapper.breaks_rings_in_transformation(pattern_to_target_map,
                                                     current):
            return False

        target_to_pattern_map = {target_atoms[matchpair.target.GetIdx()]:
                                 pattern_atoms[matchpair.pattern.GetIdx()]
                                 for matchpair in match.GetAtoms()}
        if AtomMapper.breaks_rings_in_transformation(target_to_pattern_map,
                                                     current):
            return False

        return True

    @staticmethod
    def breaks_rings_in_transformation(atom_map, current):
            """Return True if the transformation from molecule1 to molecule2 breaks rings.
    
            Parameters
            ----------
            atom_map : dict of OEAtom : OEAtom
                atom_map[molecule1_atom] is the corresponding molecule2 atom
            current : oechem.OEMol
                Initial molecule whose rings are to be checked for not being broken
    
            Returns
            -------
            bool
            """
            # not sure how this works if proposed isn't called??
            for cycle in AtomMapper.enumerate_cycle_basis(current):
                cycle_size = len(cycle)
                # first check that ALL of the ring is in the map or out
                atoms_in_cycle = set([bond.GetBgn() for bond in cycle] + [bond.GetEnd() for bond in cycle])
                number_of_cycle_atoms_mapped = 0
                for atom in atoms_in_cycle:
                    if atom in atom_map:
                        number_of_cycle_atoms_mapped += 1
                if number_of_cycle_atoms_mapped == 0:
                    # none of the ring is mapped - ALL unique, so continue
                    continue
                if number_of_cycle_atoms_mapped != len(atoms_in_cycle):
                    return True # not all atoms in ring are mapped
                # now we are sure the ring is either fully core or fully unique, check the same is true of the proposed
                for bond in cycle:
                    # then check that all of the new ring is in the map
                    if ((bond.GetBgn() in atom_map) and (bond.GetEnd() in atom_map)):
                        if not oechem.OEAtomIsInRingSize(atom_map[bond.GetBgn()], cycle_size):
                            return True
                        if not oechem.OEAtomIsInRingSize(atom_map[bond.GetEnd()], cycle_size):
                            return True
            return False  # no rings in molecule1 are broken in molecule2

    @staticmethod
    def enumerate_cycle_basis(molecule):
        """Enumerate a closed cycle basis of bonds in molecule.

        This uses cycle_basis from NetworkX:
        https://networkx.github.io/documentation/networkx-1.10/reference/generated/networkx.algorithms.cycles.cycle_basis.html#networkx.algorithms.cycles.cycle_basis

        Parameters
        ----------
        molecule : OEMol
            The molecule for a closed cycle basis of Bonds is to be identified

        Returns
        -------
        bond_cycle_basis : list of list of OEBond
            bond_cycle_basis[cycle_index] is a list of OEBond objects that define a cycle in the basis
            You can think of these as the minimal spanning set of ring systems to check.
        """
        g = nx.Graph()
        for atom in molecule.GetAtoms():
            g.add_node(atom.GetIdx())
        for bond in molecule.GetBonds():
            g.add_edge(bond.GetBgnIdx(), bond.GetEndIdx(), bond=bond)
        bond_cycle_basis = list()
        for cycle in nx.cycle_basis(g):
            bond_cycle = list()
            for i in range(len(cycle)):
                atom_index_1 = cycle[i]
                atom_index_2 = cycle[(i+1)%len(cycle)]
                edge = g.edges[atom_index_1,atom_index_2]
                bond = edge['bond']
                bond_cycle.append(bond)
            bond_cycle_basis.append(bond_cycle)
        return bond_cycle_basis

    @staticmethod
    def rank_degenerate_maps(matches, current, proposed):
        """If the atom/bond expressions for maximal substructure is relaxed,
         then the maps with the highest scores will likely be degenerate.
        Consequently, it is important to reduce the degeneracy with other tests

        This test will give each match a score wherein every atom matching
        with the same atomic number (in aromatic rings) will
        receive a +1 score.

        Parameters
        ----------
        matches : type
            Description of parameter `matches`.
        current : oechem.OEMol
            oemol of first molecule
        proposed : oechem.OEMol
            oemol of second molecule

        Returns
        -------
        list of matches
            Ordered list of the matches

        """
        """
        """
        score_list = {}
        for idx, match in enumerate(matches):
            counter_arom, counter_aliph = 0, 0
            for matchpair in match.GetAtoms():
                old_index, new_index = matchpair.pattern.GetIdx(), matchpair.target.GetIdx()
                old_atom, new_atom = current.GetAtom(oechem.OEHasAtomIdx(old_index)), proposed.GetAtom(oechem.OEHasAtomIdx(new_index))

                if old_atom.IsAromatic() and new_atom.IsAromatic(): #if both are aromatic
                    if old_atom.GetAtomicNum() == new_atom.GetAtomicNum():
                        counter_arom += 1
                else: # TODO: specify whether a single atom is aromatic/aliphatic (for ring form/break purposes)
                    old_atomic_num, new_atomic_num = old_atom.GetAtomicNum(), new_atom.GetAtomicNum()
                    if old_atomic_num != 1 and new_atomic_num == old_atomic_num:
                        counter_aliph += 1

            score_list[idx] = (counter_arom, counter_aliph)

        # return a list of matches with the most aromatic matches
        max_arom_score = max([tup[0] for tup in score_list.values()])
        top_arom_match_dict = {index: match for index, match in enumerate(matches) if score_list[index][0] == max_arom_score}

        #filter further for aliphatic matches...
        max_aliph_score = max([score_list[idx][1] for idx in top_arom_match_dict.keys()])
        top_aliph_matches = [top_arom_match_dict[idx] for idx in top_arom_match_dict.keys() if score_list[idx][1] == max_aliph_score]

        return top_aliph_matches

    def save_atom_mapping(self,filename='atom_map.png'):
        from perses.utils.smallmolecules import render_atom_mapping
        render_atom_mapping(filename, self.current_mol, self.proposed_mol, self.atom_map)

class TopologyProposal(object):
    """
    This is a container class with convenience methods to access various objects needed
    for a topology proposal

    Arguments
    ---------
    new_topology : simtk.openmm.Topology object
        openmm Topology representing the proposed new system
    new_system : simtk.openmm.System object
        openmm System of the newly proposed state
    old_topology : simtk.openmm.Topology object
        openmm Topology of the current system
    old_system : simtk.openmm.System object
        openm System of the current state
    logp_proposal : float
        contribution from the chemical proposal to the log probability of acceptance (Eq. 36 for hybrid; Eq. 53 for two-stage)
    new_to_old_atom_map : dict
        {new_atom_idx : old_atom_idx} map for the two systems
    old_alchemical_atoms : list, optional, default=None
        List of all atoms in old system that are being transformed.
        If None, all atoms are assumed to be part of the alchemical region.
    chemical_state_key : str
        The current chemical state (unique)
    metadata : dict
        additional information of interest about the state

    Properties
    ----------
    new_topology : simtk.openmm.Topology object
        openmm Topology representing the proposed new system
    new_system : simtk.openmm.System object
        openmm System of the newly proposed state
    old_topology : simtk.openmm.Topology object
        openmm Topology of the current system
    old_system : simtk.openmm.System object
        openm System of the current state
    old_positions : [n, 3] np.array, Quantity
        positions of the old system
    logp_proposal : float
        contribution from the chemical proposal to the log probability of acceptance (Eq. 36 for hybrid; Eq. 53 for two-stage)
    new_to_old_atom_map : dict
        {new_atom_idx : old_atom_idx} map for the two systems
    old_to_new_atom_map : dict
        {old_atom_idx : new_atom_idx} map for the two systems
    new_alchemical_atoms : list
        List of all atoms in new system that are being transformed
    new_environment_atoms : list
        List of all atoms in new system that are not transformed, just mapped
    old_alchemical_atoms : list
        List of all atoms in old system that are being transformed
    old_environment_atoms : list
        List of all atoms in old system that are not transformed, just mapped
    unique_new_atoms : list of int
        List of indices of the unique new atoms
    unique_old_atoms : list of int
        List of indices of the unique old atoms
    natoms_new : int
        Number of atoms in the new system
    natoms_old : int
        Number of atoms in the old system
    old_chemical_state_key : str
        The previous chemical state key
    new_chemical_state_key : str
        The proposed chemical state key
    old_residue_name : str
        Name of the old residue
    new_residue_name : str
        Name of the new residue
    metadata : dict
        additional information of interest about the state
    """

    def __init__(self,
                 new_topology=None, new_system=None,
                 old_topology=None, old_system=None,
                 logp_proposal=None,
                 new_to_old_atom_map=None, old_alchemical_atoms=None,
                 old_chemical_state_key=None, new_chemical_state_key=None,
                 old_residue_name='MOL', new_residue_name='MOL',
                 metadata=None):

        if new_chemical_state_key is None or old_chemical_state_key is None:
            raise ValueError("chemical_state_keys must be set.")
        self._new_topology = new_topology
        self._new_system = new_system
        self._old_topology = old_topology
        self._old_system = old_system
        self._logp_proposal = logp_proposal
        self._new_chemical_state_key = new_chemical_state_key
        self._old_chemical_state_key = old_chemical_state_key
        self._old_residue_name = old_residue_name
        self._new_residue_name = new_residue_name
        self._new_to_old_atom_map = new_to_old_atom_map
        self._old_to_new_atom_map = {old_atom : new_atom for new_atom, old_atom in new_to_old_atom_map.items()}
        self._unique_new_atoms = list(set(range(self._new_topology.getNumAtoms()))-set(self._new_to_old_atom_map.keys()))
        self._unique_old_atoms = list(set(range(self._old_topology.getNumAtoms()))-set(self._new_to_old_atom_map.values()))
        self._old_alchemical_atoms = set(old_alchemical_atoms) if (old_alchemical_atoms is not None) else {atom for atom in range(old_system.getNumParticles())}
        self._new_alchemical_atoms = set(self._old_to_new_atom_map.values()).union(self._unique_new_atoms)
        self._old_environment_atoms = set(range(old_system.getNumParticles())) - self._old_alchemical_atoms
        self._new_environment_atoms = set(range(new_system.getNumParticles())) - self._new_alchemical_atoms
        self._metadata = metadata

    @property
    def new_topology(self):
        return self._new_topology
    @property
    def new_system(self):
        return self._new_system
    @property
    def old_topology(self):
        return self._old_topology
    @property
    def old_system(self):
        return self._old_system
    @property
    def logp_proposal(self):
        return self._logp_proposal
    @property
    def new_to_old_atom_map(self):
        return self._new_to_old_atom_map
    @property
    def old_to_new_atom_map(self):
        return self._old_to_new_atom_map
    @property
    def unique_new_atoms(self):
        return self._unique_new_atoms
    @property
    def unique_old_atoms(self):
        return self._unique_old_atoms
    @property
    def new_alchemical_atoms(self):
        return list(self._new_alchemical_atoms)
    @property
    def old_alchemical_atoms(self):
        return list(self._old_alchemical_atoms)
    @property
    def new_environment_atoms(self):
        return list(self._new_environment_atoms)
    @property
    def old_environment_atoms(self):
        return list(self._old_environment_atoms)
    @property
    def n_atoms_new(self):
        return self._new_system.getNumParticles()
    @property
    def n_atoms_old(self):
        return self._old_system.getNumParticles()
    @property
    def new_chemical_state_key(self):
        return self._new_chemical_state_key
    @property
    def old_chemical_state_key(self):
        return self._old_chemical_state_key
    @property
    def old_residue_name(self):
        return self._old_residue_name
    @property
    def new_residue_name(self):
        return self._new_residue_name
    @property
    def metadata(self):
        return self._metadata

class ProposalEngine(object):
    """
    This defines a type which, given the requisite metadata, can produce Proposals (namedtuple)
    of new topologies.

    Arguments
    --------
    system_generator : SystemGenerator
        The SystemGenerator to use to generate new System objects for proposed Topology objects
    proposal_metadata : dict
        Contains information necessary to initialize proposal engine
    verbose : bool, optional, default=False
        If True, print verbose debugging output

    Properties
    ----------
    chemical_state_list : list of str
         a list of all the chemical states that this proposal engine may visit.
    """

    def __init__(self, system_generator, proposal_metadata=None, always_change=True, verbose=False, **kwargs):
        self._system_generator = system_generator
        self.verbose = verbose
        self._always_change = always_change
        super(ProposalEngine, self).__init__(**kwargs)

    def propose(self, current_system, current_topology, current_metadata=None):
        """
        Base interface for proposal method.

        Arguments
        ---------
        current_system : simtk.openmm.System object
            The current system object
        current_topology : simtk.openmm.app.Topology object
            The current topology
        current_metadata : dict
            Additional metadata about the state
        Returns
        -------
        proposal : TopologyProposal
            NamedTuple of type TopologyProposal containing forward and reverse
            probabilities, as well as old and new topologies and atom
            mapping
        """
        return TopologyProposal(new_topology=app.Topology(), old_topology=app.Topology(), old_system=current_system, old_chemical_state_key="C", new_chemical_state_key="C", logp_proposal=0.0, new_to_old_atom_map={0 : 0}, metadata={'molecule_smiles' : 'CC'})

    def compute_state_key(self, topology):
        """
        Compute the corresponding state key of a given topology,
        according to this proposal engine's scheme.

        Parameters
        ----------
        topology : app.Topology
            the topology in question

        Returns
        -------
        chemical_state_key : str
            The chemical_state_key
        """
        pass

    @property
    def chemical_state_list(self):
        raise NotImplementedError("This ProposalEngine does not expose a list of possible chemical states.")

class PolymerProposalEngine(ProposalEngine):
    """
    Base class for ProposalEngine implementations that modify polymer components of systems.

    This base class is not meant to be invoked directly.
    """

    # TODO: Eliminate 'verbose' option in favor of logging
    # TODO: Document meaning of 'aggregate'
    def __init__(self, system_generator, chain_id, proposal_metadata=None, verbose=False, always_change=True, aggregate=False):
        """
        Create a polymer proposal engine

        Parameters
        ----------
        system_generator : SystemGenerator
            The SystemGenerator to use to generate perturbed systems
        chain_id : str
            The chain identifier in the Topology object to be mutated
        proposal_metadata : dict, optional, default=None
            Any metadata to be maintained
        verbose : bool, optional, default=False
            If True, will generate verbose output
        always_change : bool, optional, default=True
            If True, will not propose self transitions
        aggregate : bool, optional, default=False
            ???????

        This base class is not meant to be invoked directly.
        """
        _logger.info(f"Instantiating PolymerProposalEngine")
        super(PolymerProposalEngine,self).__init__(system_generator, proposal_metadata=proposal_metadata, verbose=verbose, always_change=always_change)
        self._chain_id = chain_id # chain identifier defining polymer to be modified
        self._aminos = ['ALA', 'ARG', 'ASN', 'ASP', 'CYS', 'GLN', 'GLU', 'GLY', 'HIS', 'ILE', 'LEU', 'LYS', 'MET', 'PHE',
                        'SER', 'THR', 'TRP', 'TYR', 'VAL'] # common naturally-occurring amino acid names
                        # Note this does not include PRO since there's a problem with OpenMM's template DEBUG
        self._aggregate = aggregate # ?????????

    def propose(self, current_system, current_topology, current_metadata=None):
        """
        Generate a TopologyProposal

        Arguments
        ---------
        current_system : simtk.openmm.System object
            The current system object
        current_topology : simtk.openmm.app.Topology object
            The current topology
        current_metadata : dict -- OPTIONAL

        Returns
        -------
        proposal : TopologyProposal
            NamedTuple of type TopologyProposal containing forward and reverse
            probabilities, as well as old and new topologies and atom
            mapping
        """

        # old_topology : simtk.openmm.app.Topology
        old_topology = app.Topology()
        append_topology(old_topology, current_topology)

        # new_topology : simtk.openmm.app.Topology
        new_topology = app.Topology()
        append_topology(new_topology, current_topology)

        # Check that old_topology and old_system have same number of atoms.
        old_system = current_system
        old_topology_natoms = old_topology.getNumAtoms()  # number of topology atoms
        old_system_natoms = old_system.getNumParticles()
        if old_topology_natoms != old_system_natoms:
            msg = 'PolymerProposalEngine: old_topology has %d atoms, while old_system has %d atoms' % (old_topology_natoms, old_system_natoms)
            raise Exception(msg)

        # metadata : dict, key = 'chain_id' , value : str
        metadata = current_metadata
        if metadata is None:
            metadata = dict()
        # old_chemical_state_key : str
        old_chemical_state_key = self.compute_state_key(old_topology)

        # index_to_new_residues : dict, key : int (index) , value : str (three letter name of proposed residue)
        index_to_new_residues, metadata = self._choose_mutant(old_topology, metadata)

        # residue_map : list(tuples : simtk.openmm.app.topology.Residue (existing residue), str (three letter name of proposed residue))
        residue_map = self._generate_residue_map(old_topology, index_to_new_residues)

        for (res, new_name) in residue_map:
            if res.name == new_name:
                del(index_to_new_residues[res.index])
        if len(index_to_new_residues) == 0:
            atom_map = dict()
            for atom in old_topology.atoms():
                atom_map[atom.index] = atom.index
            if self.verbose: print('PolymerProposalEngine: No changes to topology proposed, returning old system and topology')
            topology_proposal = TopologyProposal(new_topology=old_topology, new_system=old_system, old_topology=old_topology, old_system=old_system, old_chemical_state_key=old_chemical_state_key, new_chemical_state_key=old_chemical_state_key, logp_proposal=0.0, new_to_old_atom_map=atom_map)
            return topology_proposal

        elif len(index_to_new_residues) > 1:
            raise Exception("Attempting to mutate more than one residue at once: ", index_to_new_residues, " The geometry engine cannot handle this.")

        # Add modified_aa property to residues in old topology
        for res in old_topology.residues():
            res.modified_aa = True if res.index in index_to_new_residues.keys() else False

        # Identify differences between old topology and proposed changes
        # excess_atoms : list(simtk.openmm.app.topology.Atom) atoms from existing residue not in new residue
        # excess_bonds : list(tuple (simtk.openmm.app.topology.Atom, simtk.openmm.app.topology.Atom)) bonds from existing residue not in new residue
        # missing_bonds : list(tuple (simtk.openmm.app.topology._TemplateAtomData, simtk.openmm.app.topology._TemplateAtomData)) bonds from new residue not in existing residue
        excess_atoms, excess_bonds, missing_atoms, missing_bonds = self._identify_differences(old_topology, residue_map)

        # Delete excess atoms and bonds from old topology
        excess_atoms_bonds = excess_atoms + excess_bonds
        new_topology = self._delete_atoms(old_topology, excess_atoms_bonds)

        # Add missing atoms and bonds to new topology
        new_topology = self._add_new_atoms(new_topology, missing_atoms, missing_bonds, residue_map)

        # index_to_new_residues : dict, key : int (index) , value : str (three letter name of proposed residue)
        atom_map = self._construct_atom_map(residue_map, old_topology, index_to_new_residues, new_topology)

        # new_chemical_state_key : str
        new_chemical_state_key = self.compute_state_key(new_topology)
        # new_system : simtk.openmm.System

        # Copy periodic box vectors from current topology
        new_topology.setPeriodicBoxVectors(current_topology.getPeriodicBoxVectors())

        # Build system
        # TODO: Remove build_system() branch once we convert entirely to new openmm-forcefields SystemBuilder
        if hasattr(self._system_generator, 'create_system'):
            new_system = self._system_generator.create_system(new_topology)
        else:
            new_system = self._system_generator.build_system(new_topology)

        # Adjust logp_propose based on HIS presence
        his_residues = ['HID', 'HIE']
        old_residue = residue_map[0][0]
        proposed_residue = residue_map[0][1]
        if old_residue.name in his_residues and proposed_residue not in his_residues:
            logp_propose = math.log(2)
        elif old_residue.name not in his_residues and proposed_residue in his_residues:
            logp_propose = math.log(0.5)
        else:
            logp_propose = 0.0

        # Create TopologyProposal.
        topology_proposal = TopologyProposal(new_topology=new_topology, new_system=new_system, old_topology=old_topology, old_system=old_system, old_chemical_state_key=old_chemical_state_key, new_chemical_state_key=new_chemical_state_key, logp_proposal=logp_propose, new_to_old_atom_map=atom_map)

        # Check that old_topology and old_system have same number of atoms.
        old_topology_natoms = old_topology.getNumAtoms()  # number of topology atoms
        old_system_natoms = old_system.getNumParticles()
        if old_topology_natoms != old_system_natoms:
            msg = 'PolymerProposalEngine: old_topology has %d atoms, while old_system has %d atoms' % (old_topology_natoms, old_system_natoms)
            raise Exception(msg)

        # Check that new_topology and new_system have same number of atoms.
        new_topology_natoms = new_topology.getNumAtoms()  # number of topology atoms
        new_system_natoms = new_system.getNumParticles()
        if new_topology_natoms != new_system_natoms:
            msg = 'PolymerProposalEngine: new_topology has %d atoms, while new_system has %d atoms' % (new_topology_natoms, new_system_natoms)
            raise Exception(msg)

        # Check to make sure no out-of-bounds atoms are present in new_to_old_atom_map
        natoms_old = topology_proposal.old_system.getNumParticles()
        natoms_new = topology_proposal.new_system.getNumParticles()
        if not set(topology_proposal.new_to_old_atom_map.values()).issubset(range(natoms_old)):
            msg = "Some new atoms in TopologyProposal.new_to_old_atom_map are not in span of new atoms (1..%d):\n" % natoms_new
            msg += str(topology_proposal.new_to_old_atom_map)
            raise Exception(msg)
        if not set(topology_proposal.new_to_old_atom_map.keys()).issubset(range(natoms_new)):
            msg = "Some new atoms in TopologyProposal.new_to_old_atom_map are not in span of old atoms (1..%d):\n" % natoms_new
            msg += str(topology_proposal.new_to_old_atom_map)
            raise Exception(msg)

        return topology_proposal

    def _choose_mutant(self, topology, metadata):
        index_to_new_residues = dict()
        return index_to_new_residues, metadata

    def _generate_residue_map(self, topology, index_to_new_residues):
        """
        generates list to reference residue instance to be edited, because topology.residues() cannot be referenced directly by index

        Arguments
        ---------
        topology : simtk.openmm.app.Topology
        index_to_new_residues : dict
            key : int (index, zero-indexed in chain)
            value : str (three letter residue name)
        Returns
        -------
        residue_map : list(tuples)
            simtk.openmm.app.topology.Residue (existing residue), str (three letter residue name of proposed residue)
        """
        # residue_map : list(tuples : simtk.openmm.app.topology.Residue (existing residue), str (three letter residue name of proposed residue))
        # r : simtk.openmm.app.topology.Residue, r.index : int, 0-indexed
        residue_map = [(r, index_to_new_residues[r.index]) for r in topology.residues() if r.index in index_to_new_residues]
        return residue_map

    def _identify_differences(self, topology, residue_map):
        """
        Identify excess atoms, excess bonds, missing atoms, and missing bonds.

        Arguments
        ---------
        topology : simtk.openmm.app.Topology
            The original Topology object to be processed
        residue_map : list(tuples)
            simtk.openmm.app.topology.Residue (existing residue), str (three letter residue name of proposed residue)
        Returns
        -------
        excess_atoms : list(simtk.openmm.app.topology.Atom)
            atoms from existing residue not in new residue
        excess_bonds : list(tuple (simtk.openmm.app.topology.Atom, simtk.openmm.app.topology.Atom))
            bonds from existing residue not in new residue
        missing_bonds : list(tuple (simtk.openmm.app.topology._TemplateAtomData, simtk.openmm.app.topology._TemplateAtomData))
            bonds from new residue not in existing residue
        """
        # excess_atoms : list(simtk.openmm.app.topology.Atom) atoms from existing residue not in new residue
        excess_atoms = list()
        # excess_bonds : list(tuple (simtk.openmm.app.topology.Atom, simtk.openmm.app.topology.Atom)) bonds from existing residue not in new residue
        excess_bonds = list()
        # missing_atoms : dict, key : simtk.openmm.app.topology.Residue, value : list(simtk.openmm.app.topology._TemplateAtomData)
        missing_atoms = dict()
        # missing_bonds : list(tuple (simtk.openmm.app.topology._TemplateAtomData, simtk.openmm.app.topology._TemplateAtomData)) bonds from new residue not in existing residue
        missing_bonds = list()

        # residue : simtk.openmm.app.topology.Residue (existing residue)
        for k, (residue, replace_with) in enumerate(residue_map):
            # Load residue template for residue to replace with
            template = self._templates[replace_with]

            # template_atom_names : dict, key : template atom index, value : template atom name
            template_atom_names = {}
            for atom in template.atoms:
                template_atom_names[template.getAtomIndexByName(atom.name)] = atom.name

            # template_atoms : list(simtk.openmm.app.topology._TemplateAtomData) atoms in new residue
            template_atoms = list(template.atoms)


            # template_bonds : dict, key : simtk.openmm.app.topology.Atom, value : str template atom name
            template_bonds = [(template_atom_names[bond[0]], template_atom_names[bond[1]]) for bond in template.bonds]
            # old_atom_names : set of unique atom names within existing residue : str
            old_atom_names = set(atom.name for atom in residue.atoms())

            # Make a list of atoms in the existing residue that are not in the new residue
            # atom : simtk.openmm.app.topology.Atom in existing residue
            for atom in residue.atoms():
                if atom.name not in template_atom_names.values():
                    excess_atoms.append(atom)

            # if residue == chain_residues[0]: # this doesn't apply?
            # template_atoms = [atom for atom in template_atoms if atom.name not in ('P', 'OP1', 'OP2')]

            # Make a list of atoms in the new residue that are not in the existing residue
            # missing : list(simtk.openmm.app.topology._TemplateAtomData) atoms in new residue not found in existing residue
            missing = list()
            # atom : simtk.openmm.app.topology._TemplateAtomData atoms in new residue
            for atom in template_atoms:
                if atom.name not in old_atom_names:
                    missing.append(atom)

            if len(missing) > 0:
                missing_atoms[residue] = missing
            else:
                missing_atoms[residue] = []

            # Make a dictionary to map atom names in old residue to atom object
            # old_atom_map : dict, key : str (atom name) , value : simtk.openmm.app.topology.Atom
            old_atom_map = dict()
            # atom : simtk.openmm.app.topology.Atom
            for atom in residue.atoms():
                # atom.name : str
                old_atom_map[atom.name] = atom

            # Make a dictionary to map atom names in new residue to atom object
            # new_atom_map : dict, key : str (atom name) , value : simtk.openmm.app.topology.Atom
            new_atom_map = dict()
            # atom : simtk.openmm.app.topology.Atom
            for atom in template_atoms:
                # atom.name : str
                new_atom_map[atom.name] = atom

            # Make a list of bonds already existing in new residue
            # old_bonds : list(tuple(str (atom name), str (atom name))) bonds between atoms both within old residue
            old_bonds = list()
            # bond : tuple(simtk.openmm.app.topology.Atom, simtk.openmm.app.topology.Atom)
            for bond in topology.bonds():
                if bond[0].residue == residue and bond[1].residue == residue:
                    old_bonds.append((bond[0].name, bond[1].name))

            # Add any bonds that exist in old residue but not in template to excess_bonds
            for bond in old_bonds:
                if bond not in template_bonds and (bond[1], bond[0]) not in template_bonds:
                    excess_bonds.append((old_atom_map[bond[0]], old_atom_map[bond[1]]))

            # Add any bonds that exist in template but not in old residue to missing_bonds
            for bond in template_bonds:
                if bond not in old_bonds and (bond[1], bond[0]) not in old_bonds:
                    missing_bonds.append((new_atom_map[bond[0]], new_atom_map[bond[1]]))

        return(excess_atoms, excess_bonds, missing_atoms, missing_bonds)

    def _delete_atoms(self, topology, to_delete):
        """
        Delete excess atoms (and corresponding bonds) from specified topology

        Parameters
        ----------
        topology : simtk.openmm.app.Topology
            The Topology to be processed
        excess_atoms : list(simtk.openmm.app.topology.Atom)
            atoms from existing residue not in new residue
        excess_bonds : list(tuple (simtk.openmm.app.topology.Atom, simtk.openmm.app.topology.Atom))
            bonds from old residue not in new residue

        Returns
        -------
        topology : simtk.openmm.app.Topology
            extra atoms and bonds from old residue have been deleted, missing atoms and bottoms in new residue not yet added

        """
        new_topology = app.Topology()
        new_topology.setPeriodicBoxVectors(topology.getPeriodicBoxVectors())

        # new_atoms : dict, key : simtk.openmm.app.topology.Atom, value : simtk.openmm.app.topology.Atom maps old atoms to the corresponding Atom in the new residue
        new_atoms = {}
        delete_set = set(to_delete)

        for chain in topology.chains():
            if chain not in delete_set:
                new_chain = new_topology.addChain(chain.id)
                for residue in chain.residues():
                    new_residue = new_topology.addResidue(residue.name, new_chain, residue.id)
                    for atom in residue.atoms():
                        if atom not in delete_set:
                            new_atom = new_topology.addAtom(atom.name, atom.element, new_residue, atom.id)
                            new_atom.old_index = atom.index
                            new_atoms[atom] = new_atom
        for bond in topology.bonds():
            if bond[0] in new_atoms and bond[1] in new_atoms:
                if bond not in delete_set and (bond[1], bond[0]) not in delete_set:
                    new_topology.addBond(new_atoms[bond[0]], new_atoms[bond[1]])
        return new_topology


    def _add_new_atoms(self, topology, missing_atoms, missing_bonds, residue_map):
        """
        Add new atoms (and corresponding bonds) to new residues
        Arguments
        ---------
        topology : simtk.openmm.app.Topology
            extra atoms from old residue have been deleted, missing atoms in new residue not yet added
        missing_atoms : dict
            key : simtk.openmm.app.topology.Residue
            value : list(simtk.openmm.app.topology._TemplateAtomData)
        missing_bonds : list(tuple (simtk.openmm.app.topology._TemplateAtomData, simtk.openmm.app.topology._TemplateAtomData))
            bonds from new residue not in existing residue
        residue_map : list(tuples)
            simtk.openmm.app.topology.Residue, str (three letter residue name of new residue)
        Returns
        -------
        topology : simtk.openmm.app.Topology
            new residues have all correct atoms and bonds for desired mutation
        """

        new_topology = app.Topology()
        new_topology.setPeriodicBoxVectors(topology.getPeriodicBoxVectors())
        # new_atoms : dict, key : simtk.openmm.app.topology.Atom, value : simtk.openmm.app.topology.Atom maps old atoms to the corresponding Atom in the new residue
        new_atoms = {}
        # new_atom_names : dict, key : str new atom name, value : simtk.openmm.app.topology.Atom maps name of new atom to the corresponding Atom in the new residue (only contains map for missing residue)
        new_atom_names = {}
        # old_residues : list(simtk.openmm.app.topology.Residue)
        old_residues = [old.index for old, new in residue_map]
        for chain in topology.chains():
            new_chain = new_topology.addChain(chain.id)
            for residue in chain.residues():
                new_residue = new_topology.addResidue(residue.name, new_chain, residue.id)
                # Add modified property to residues in new topology
                new_residue.modified_aa = True if residue.index in old_residues else False
                # Copy over atoms from old residue to new residue
                for atom in residue.atoms():
                    # new_atom : simtk.openmm.app.topology.Atom
                    new_atom = new_topology.addAtom(atom.name, atom.element, new_residue)
                    new_atom.old_index = atom.old_index
                    new_atoms[atom] = new_atom
                    if new_residue.modified_aa:
                        new_atom_names[new_atom.name] = new_atom
                # Check if old residue is in residue_map
                # old_residue : simtk.openmm.app.topology.Residue (old residue)
                # new_residue_name : str (three letter residue name of new residue)
                for i, (old_residue, new_residue_name) in enumerate(residue_map):
                    if self._is_residue_equal(residue, old_residue):
                        # Add missing atoms to new residue
                        # atom : simtk.openmm.app.topology._TemplateAtomData
                        for atom in missing_atoms[old_residue]:
                            new_atom = new_topology.addAtom(atom.name, atom.element, new_residue)
                            new_atoms[atom] = new_atom
                            new_atom_names[new_atom.name] = new_atom
                        new_residue.name = residue_map[i][1]

        # Copy over bonds from topology to new topology
        for bond in topology.bonds():
            new_topology.addBond(new_atoms[bond[0]], new_atoms[bond[1]])

        for bond in missing_bonds:
            new_topology.addBond(new_atom_names[bond[0].name], new_atom_names[bond[1].name])

        return new_topology

    def _is_residue_equal(self, residue, other_residue):
        """
            Check if residue is equal to other_residue based on their names, indices, ids, and chain ids.
            Arguments
            ---------
            residue : simtk.openmm.app.topology.Residue
            other_residue : simtk.openmm.app.topology.Residue
            Returns
            -------
            boolean True if residues are equal, otherwise False
        """
        return residue.name == other_residue.name and residue.index == other_residue.index and residue.chain.id == other_residue.chain.id and residue.id == other_residue.id

    def _construct_atom_map(self, residue_map, old_topology, index_to_new_residues, new_topology):
        """
        Construct atom map (key: index to new residue, value: index to old residue) to supply as an argument to the TopologyProposal.
        Arguments
        ---------
        residue_map : list(tuples)
            simtk.openmm.app.topology.Residue, str (three letter residue name of new residue)
        old_topology : simtk.openmm.app.Topology
        index_to_new_residues : dict, key : int (index) , value : str (three letter name of proposed residue)
        new_topology : simtk.openmm.app.Topology
        Returns
        -------
        atom_map : dict, key: int (index
            new residues have all correct atoms and bonds for desired mutation
        """

        # atom_map : dict, key : int (index of atom in old topology) , value : int (index of same atom in new topology)
        atom_map = dict()

        # atoms with an old_index attribute should be mapped
        # k : int
        # atom : simtk.openmm.app.topology.Atom
        def match_backbone(old_residue, new_residue, atom_name):
            """
            Forcibly including CA and N in the map even if they don't meet
            matching criteria
            """
            found_old_atom = False
            for atom in old_residue.atoms():
                if atom.name == atom_name:
                    old_atom = atom
                    found_old_atom = True
                    break
            assert found_old_atom
            found_new_atom = False
            for atom in new_residue.atoms():
                if atom.name == atom_name:
                    new_atom = atom
                    found_new_atom = True
                    break
            assert found_new_atom
            return new_atom.index, old_atom.index

        # old_to_new_residues : dict, key : str old residue name, key : simtk.openmm.app.topology.Residue new residue
        old_to_new_residues = {}
        for old_residue in old_topology.residues():
            for new_residue in new_topology.residues():
                if old_residue.index == new_residue.index:
                    old_to_new_residues[old_residue.name] = new_residue
                    break

        # modified_residues : dict, key : index of old residue, value : proposed residue
        modified_residues = dict()
        for map_entry in residue_map:
            old_residue = map_entry[0]
            modified_residues[old_residue.index] = old_to_new_residues[old_residue.name]

        # old_residues : dict, key : index of old residue, value : old residue
        old_residues = dict()
        for residue in old_topology.residues():
            if residue.index in index_to_new_residues.keys():
                old_residues[residue.index] = residue

        # Create initial atom map for atoms in new topology that are not part of modified residues
        for atom in new_topology.atoms():
            if atom.residue in modified_residues.values():
                continue
            try:
                atom_map[atom.index] = atom.old_index
            except AttributeError:
                pass

        # Update atom map with atom mappings for residues that have been modified
        for index in index_to_new_residues.keys():
            old_res = old_residues[index]
            new_res = modified_residues[index]

            # Save index of first atom in old residue and new residue
            for atom in old_res.atoms():
                first_atom_index_old = atom.index
                break
            for atom in new_res.atoms():
                first_atom_index_new = atom.index
                break

            old_oemol_res = FFAllAngleGeometryEngine._oemol_from_residue(old_res)
            new_oemol_res = FFAllAngleGeometryEngine._oemol_from_residue(new_res)
            # local_atom_map : dict, key : index of atom in new residue, value : index of atom in old residue.
            local_atom_map = self._get_mol_atom_matches(old_oemol_res, new_oemol_res, first_atom_index_old, first_atom_index_new)
            for backbone_name in ['CA','N']:
                new_index, old_index = match_backbone(old_residues[index], modified_residues[index], backbone_name)
                local_atom_map[new_index] = old_index
            atom_map.update(local_atom_map)
        return atom_map

    def _get_mol_atom_matches(self, current_molecule, proposed_molecule, first_atom_index_old, first_atom_index_new):
        """
        Given two molecules, returns the mapping of atoms between them.
        Arguments
        ---------
        current_molecule : openeye.oechem.oemol object
             The current molecule in the sampler
        proposed_molecule : openeye.oechem.oemol object
             The proposed new molecule
        first_atom_index_old : int
            The index of the first atom in the old resiude/current molecule
        first_atom_index_new : int
            The index of the first atom in the new residue/proposed molecule
        Note: Since FFAllAngleGeometryEngine._oemol_from_residue creates a new topology for the specified residue,
        the atom indices in the output oemol (i.e. current_molecule and proposed_molecule) are reset to start at 0.
        Therefore, first_atom_index_old and first_atom_index_new are used to correct the indices such that they match
        the atom indices of the original old and new residues.
        Returns
        -------
        new_to_old_atom_map : dict, key : index of atom in new residue, value : index of atom in old residue
        """
        # Load current and proposed residues as OEGraphMol objects
        oegraphmol_current = oechem.OEGraphMol(current_molecule)
        oegraphmol_proposed = oechem.OEGraphMol(proposed_molecule)

        # Instantiate Maximum Common Substructures (MCS) object and intialize properties
        mcs = oechem.OEMCSSearch(oechem.OEMCSType_Exhaustive)
        atomexpr = oechem.OEExprOpts_Aromaticity | oechem.OEExprOpts_RingMember | oechem.OEExprOpts_Degree | oechem.OEExprOpts_AtomicNumber
        bondexpr = oechem.OEExprOpts_Aromaticity | oechem.OEExprOpts_RingMember
        mcs.Init(oegraphmol_current, atomexpr, bondexpr)

        def forcibly_matched(mcs, proposed, atom_name):
            old_atom = list(mcs.GetPattern().GetAtoms(atom_name))
            assert len(old_atom) == 1
            old_atom = old_atom[0]

            new_atom = list(proposed.GetAtoms(atom_name))
            assert len(new_atom) == 1
            new_atom = new_atom[0]
            return old_atom, new_atom

        # Forcibly match C and O atoms in proposed residue to atoms in current residue
        for matched_atom_name in ['C','O']:
            force_match = oechem.OEHasAtomName(matched_atom_name)
            old_atom, new_atom = forcibly_matched(mcs, oegraphmol_proposed, force_match)
            this_match = oechem.OEMatchPairAtom(old_atom, new_atom)
            assert mcs.AddConstraint(this_match)

        # Generate matches using MCS
        mcs.SetMCSFunc(oechem.OEMCSMaxBondsCompleteCycles())
        unique = True
        matches = [m for m in mcs.Match(oegraphmol_proposed, unique)]

        # Handle case where there are no matches
        if len(matches) == 0:
            from perses.utils.openeye import describe_oemol
            msg = 'No matches found in _get_mol_atom_matches.\n'
            msg += '\n'
            msg += 'oegraphmol_current:\n'
            msg += describe_oemol(oegraphmol_current)
            msg += '\n'
            msg += 'oegraphmol_proposed:\n'
            msg += describe_oemol(oegraphmol_proposed)
            raise Exception(msg)

        # Select match and generate atom map
        match = np.random.choice(matches)
        new_to_old_atom_map = {}
        for match_pair in match.GetAtoms():
            if match_pair.pattern.GetAtomicNum() == 1 and match_pair.target.GetAtomicNum() == 1:  # Do not map hydrogens
                continue
            O2_index_current = current_molecule.NumAtoms() - 2
            O2_index_proposed = proposed_molecule.NumAtoms() -2
            if 'O2' in match_pair.pattern.GetName() and 'O2' in match_pair.target.GetName() and match_pair.pattern.GetIdx() == O2_index_current and match_pair.target.GetIdx() == O2_index_proposed:  # Do not map O2 if its second to last index in atom (this O2 was added to oemol to complete the residue)
                continue
            old_index = match_pair.pattern.GetData("topology_index")
            new_index = match_pair.target.GetData("topology_index")
            new_to_old_atom_map[new_index] = old_index

        return new_to_old_atom_map

    def compute_state_key(self, topology):
        for chain in topology.chains():
            if chain.id == self._chain_id:
                break
        chemical_state_key = ''
        for index, res in enumerate(chain.residues()):
            if index > 0:
                chemical_state_key += '-'
            chemical_state_key += res.name

        return chemical_state_key

class PointMutationEngine(PolymerProposalEngine):
    """
    ProposalEngine for generating point mutation variants of a wild-type polymer

    Examples
    --------

    Mutations of a terminally-blocked peptide

    >>> from openmmtools.testsystems import AlanineDipeptideExplicit
    >>> testsystem = AlanineDipeptideExplicit()
    >>> system, topology, positions = testsystem.system, testsystem.topology, testsystem.positions

    >>> from topology_proposal import PointMutationEngine
    >>> engine = PointMutationEngine(topology, system_generator, chain_id='A', residues_allowed_to_mutate='max_point_mutants=1)

    """

    # TODO: Overhaul API to make it easier to specify mutations
    def __init__(self, wildtype_topology, system_generator, chain_id, proposal_metadata=None, max_point_mutants=1, residues_allowed_to_mutate=None, allowed_mutations=None, verbose=False, always_change=True, aggregate=False):
        """
        Create a PointMutationEngine for proposing point mutations of a biopolymer component of a system.

        Note: When instantiating this class, be sure to include cap residues ('ACE', 'NME') in the input topology.
        Otherwise, mutating the first residue will result in a residue template mismatch.
        See openmmtools/data/alanine-dipeptide-gbsa/alanine-dipeptide.pdb for reference.

        Parameters
        ----------
        wildtype_topology : openmm.app.Topology
            The Topology object describing the system
        system_generator : SystemGenerator
            The SystemGenerator to generate new parameterized System objects
        chain_id : str
            id of the chain to mutate
            (using the first chain with the id, if there are multiple)
        proposal_metadata : dict, optional, default=None
            Contains information necessary to initialize proposal engine
        max_point_mutants : int, optional, default=1
            If not None, limit the number of point mutations that are allowed simultaneously
        residues_allowed_to_mutate : list(str), optional, default=None
            Contains residue ids
            If not specified, engine assumes all residues (except ACE and NME caps) may be mutated.
        allowed_mutations : list(tuple), optional, default=None
            ('residue id to mutate','desired mutant residue name (3-letter code)')
            Example:
                Desired systems are wild type T4 lysozyme, T4 lysozyme L99A, and T4 lysozyme L99A/M102Q
                allowed_mutations = [
                    ('99', 'ALA'),
                    ('102','GLN')
                ]
            If this is not specified, the engine will propose a random amino acid at a random location.
        always_change : bool, optional, default=True
            Have the proposal engine always propose a state different from the current state.
            If the current state is WT, always propose a mutation.
            If the current state is mutant, always propose a different mutant or WT.
        aggregate : bool, optional, default=False
            Have the proposal engine aggregate mutations.
            If aggregate is set to False, the engine will undo mutants in the current topology such that the next proposal
            if the WT state
            If aggregate is set to True, the engine will not undo the mutants from the current topology, thereby allowing
            each proposal to contain multiple mutations.

        """
        super(PointMutationEngine,self).__init__(system_generator, chain_id, proposal_metadata=proposal_metadata, verbose=verbose, always_change=always_change, aggregate=aggregate)

        assert isinstance(wildtype_topology, app.Topology)

        # Check that provided topology has specified chain.
        chain_ids_in_topology = [chain.id for chain in wildtype_topology.chains()]
        if chain_id not in chain_ids_in_topology:
            raise Exception("Specified chain_id '%s' not found in provided wildtype_topology. Choices are: %s" % (chain_id, str(chain_ids_in_topology)))

        if max_point_mutants != 1:
            raise ValueError('max_point_mutants != 1 not yet supported')

        self._max_point_mutants = max_point_mutants
        self._wildtype = wildtype_topology
        self._ff = system_generator.forcefield
        self._templates = self._ff._templates
        self._residues_allowed_to_mutate = residues_allowed_to_mutate
        self._allowed_mutations = allowed_mutations
        if proposal_metadata is None:
            proposal_metadata = dict()
        self._metadata = proposal_metadata

    def _choose_mutant(self, topology, metadata):
        chain_id = self._chain_id
        old_key = self._compute_mutant_key(topology, chain_id)
        index_to_new_residues = self._undo_old_mutants(topology, chain_id, old_key)
        if index_to_new_residues and not self._aggregate:  # Starting state is mutant and mutations should not be aggregated
            pass  # At this point, index_to_new_residues contains mutations that need to be undone. This iteration will result in WT,
        else:  # Starting state is WT or starting state is mutant and mutations should be aggregated
            index_to_new_residues = dict()
            if self._allowed_mutations is not None:
                allowed_mutations = self._allowed_mutations
                index_to_new_residues = self._choose_mutation_from_allowed(topology, chain_id, allowed_mutations, index_to_new_residues, old_key)
            else:
                # index_to_new_residues : dict, key : int (index) , value : str (three letter residue name)
                index_to_new_residues = self._propose_mutation(topology, chain_id, index_to_new_residues)
        # metadata['mutations'] : list(str (three letter WT residue name - index - three letter MUT residue name) )
        metadata['mutations'] = self._save_mutations(topology, index_to_new_residues)
        return index_to_new_residues, metadata

    def _undo_old_mutants(self, topology, chain_id, old_key):
        index_to_new_residues = dict()
        if old_key == 'WT':
            return index_to_new_residues
        for chain in topology.chains():
            if chain.id == chain_id:
                break
        residue_id_to_index = {residue.id : residue.index for residue in chain.residues()}
        for mutant in old_key.split('-'):
            old_res = mutant[:3]
            residue_id = mutant[3:-3]
            index_to_new_residues[residue_id_to_index[residue_id]] = old_res
        return index_to_new_residues

    def _choose_mutation_from_allowed(self, topology, chain_id, allowed_mutations, index_to_new_residues, old_key):
        """
        Used when allowed mutations have been specified
        Assume (for now) uniform probability of selecting each specified mutant
        Arguments
        ---------
        topology : simtk.openmm.app.Topology
        chain_id : str
        allowed_mutations : list(tuple)
            list of allowed mutations -- each mutation is a tuple of the residue id and three-letter amino acid code of desired mutant
        index_to_new_residues : dict
            key : int (index, zero-indexed in chain)
            value : str (three letter residue name)
            contains information to mutate back to WT as starting point for new mutants
        old_key : str
            chemical_state_key for old topology
        Returns
        -------
        index_to_new_residues : dict
            key : int (index, zero-indexed in chain)
            value : str (three letter residue name)
        """

        # Set chain and create id-index mapping for residues in chain
        chain_found = False
        for anychain in topology.chains():
            if anychain.id == chain_id:
                # chain : simtk.openmm.app.topology.Chain
                chain = anychain
                chain_found = True
                break
        if not chain_found:
            chains = [chain.id for chain in topology.chains()]
            raise Exception("Chain '%s' not found in Topology. Chains present are: %s" % (chain_id, str(chains)))
        residue_id_to_index = {residue.id : residue.index for residue in chain.residues()}

        # Define location probabilities and propose a location/mutant state
        if self._always_change:
            # location_prob : np.array, probability value for each mutant state at their respective locations in allowed_mutations (uniform).
            if old_key == 'WT':
                location_prob = [1.0 / len(allowed_mutations)] * len(allowed_mutations)
                proposed_location = np.random.choice(range(len(allowed_mutations)), p=location_prob)
            else:
                current_mutations = []
                for mutant in old_key.split('-'):
                    residue_id = mutant[3:-3]
                    new_res = mutant[-3:]
                    current_mutations.append((residue_id, new_res))

                new_mutations = []
                for mutation in allowed_mutations:
                    if mutation not in current_mutations:
                        new_mutations.append(mutation)
                if not new_mutations:
                    raise Exception("The old topology state contains all allowed mutations (%s). Please specify additional mutations." % allowed_mutations[0])

                location_prob = [1.0 / (len(new_mutations))] * len(new_mutations)
                proposed_location = np.random.choice(range(len(new_mutations)), p=location_prob)

        else:
            if old_key == 'WT':
                location_prob = [1.0 / (len(allowed_mutations)+1.0)] * (len(allowed_mutations)+1)
                proposed_location = np.random.choice(range(len(allowed_mutations) + 1), p=location_prob)
            else:
                location_prob = [1.0 / len(allowed_mutations)] * len(allowed_mutations)
                proposed_location = np.random.choice(range(len(allowed_mutations)), p=location_prob)

        # If the proposed state is the same as the current state
        # index_to_new_residues : dict, key : int (index of residue, 0-indexed), value : str (three letter residue name)
        if old_key == 'WT' and proposed_location == len(allowed_mutations):
            # Choose WT
            return index_to_new_residues
        elif old_key != 'WT':
            for mutant in old_key.split('-'):
                residue_id = mutant[3:-3]
                new_res = mutant[-3:]
                if allowed_mutations.index((residue_id, new_res)) == proposed_location:
                    return index_to_new_residues

        residue_id = allowed_mutations[proposed_location][0]
        residue_name = allowed_mutations[proposed_location][1]
        # Verify residue with mutation exists in old topology and is not the first or last residue
        # original_residue : simtk.openmm.app.topology.Residue
        original_residue = ''
        for res in chain.residues():
            if res.index == residue_id_to_index[residue_id]:
                original_residue = res
                break
        if not original_residue:
            raise Exception("User-specified an allowed mutation at residue %s , but that residue does not exist" % residue_id)
        if original_residue.index == 0 or original_residue.index == topology.getNumResidues() - 1:
            raise Exception("Residue not found. Be sure you are not trying to mutate the first or last residue."
                            " If you wish to modify one of these residues, make sure you have added cap residues to the input topology.")

        # Check if mutated residue's name is same as residue's name in old topology
        if original_residue.name in ['HID', 'HIE']:
            original_residue.name = 'HIS'
        if original_residue.name == residue_name:
            return index_to_new_residues

        # Save proposed mutation to index_to_new_residues
        # index_to_new_residues : dict, key : int (index of residue, 0-indexed), value : str (three letter residue name)
        index_to_new_residues[residue_id_to_index[residue_id]] = residue_name

        # Randomly choose HIS template ('HIS' does not exist as a template)
        if residue_name == 'HIS':
            his_state = ['HIE','HID']
            his_prob = [1/len(his_state)] * len(his_state)
            his_choice = np.random.choice(range(len(his_state)), p=his_prob)
            index_to_new_residues[residue_id_to_index[residue_id]] = his_state[his_choice]
        return index_to_new_residues

    def _propose_mutation(self, topology, chain_id, index_to_new_residues):
        """
        Arguments
        ---------
        topology : simtk.openmm.app.Topology
        chain_id : str
        index_to_new_residues : dict
            key : int (index, zero-indexed in chain)
            value : str (three letter residue name)
            contains information to mutate back to WT as starting point for new mutants
        old_key : str
            chemical_state_key for old topology
        Returns
        -------
        index_to_new_residues : dict
            key : int (index, zero-indexed in chain)
            value : str (three letter residue name)
        """

        # Set chain and create id-index mapping for residues in chain
        # chain : simtk.openmm.app.topology.Chain
        chain_found = False
        for anychain in topology.chains():
            if anychain.id == chain_id:
                chain = anychain
                chain_found = True
                residue_id_to_index = {residue.id: residue.index for residue in chain.residues()}
                if self._residues_allowed_to_mutate is None:
                    chain_residues = [res for res in chain.residues() if res.index != 0 and res.index != topology.getNumResidues()-1]
                    # num_residues : int
                    num_residues = len(chain_residues)
                else:
                    for res_id in self._residues_allowed_to_mutate:
                        if res_id not in residue_id_to_index.keys():
                            raise Exception(
                                "Residue id '%s' not found in Topology. Residue ids present are: %s. "
                                "\n\t Note: The type of the residue id must be 'str'" % (res_id, str(residue_id_to_index.keys())))
                    num_residues = len(self._residues_allowed_to_mutate)
                    chain_residues = self._mutable_residues(chain)
                break
        if not chain_found:
            chains = [chain.id for chain in topology.chains()]
            raise Exception("Chain '%s' not found in Topology. Chains present are: %s" % (chain_id, str(chains)))

        # Define location probabilities
        # location_prob : np.array, probability value for each residue location (uniform)
        location_prob = [1.0/num_residues] * num_residues

        # Propose a location at which to mutate the residue
        # proposed_location : int, index of chosen entry in location_prob
        proposed_location = np.random.choice(range(num_residues), p=location_prob)

        # Rename residue to HIS if it uses one of the HIS-derived templates
        # original_residue : simtk.openmm.app.topology.Residue
        original_residue = chain_residues[proposed_location]
        if original_residue.name in ['HIE', 'HID']:
            original_residue.name = 'HIS'

        if self._residues_allowed_to_mutate is None:
            proposed_location = original_residue.index
        else:
            proposed_location = residue_id_to_index[self._residues_allowed_to_mutate[proposed_location]]

        # Define probabilities for amino acid options and choose one
        # amino_prob : np.array, probability value for each amino acid option (uniform)
        aminos = self._aminos
        if self._always_change:
            amino_prob = [1.0/(len(aminos)-1)] * len(aminos)
            amino_prob[aminos.index(original_residue.name)] = 0.0
        else:
            amino_prob = [1.0/len(aminos)] * len(aminos)
        # proposed_amino_index : int, index of three letter residue name in aminos list
        proposed_amino_index = np.random.choice(range(len(aminos)), p=amino_prob)

        # Save proposed mutation to index_to_new_residues
        # index_to_new_residues : dict, key : int (index of residue, 0-indexed), value : str (three letter residue name)
        index_to_new_residues[proposed_location] = aminos[proposed_amino_index]

        # Randomly choose HIS template ('HIS' does not exist as a template)
        if aminos[proposed_amino_index] == 'HIS':
            his_state = ['HIE','HID']
            his_prob = [1 / len(his_state)] * len(his_state)
            his_choice = np.random.choice(range(len(his_state)), p=his_prob)
            index_to_new_residues[proposed_location] = his_state[his_choice]
        return index_to_new_residues

    def _mutable_residues(self, chain):
        chain_residues = [residue for residue in chain.residues() if residue.id in self._residues_allowed_to_mutate]
        return chain_residues

    def _save_mutations(self, topology, index_to_new_residues):
        """
        Arguments
        ---------
        topology : simtk.openmm.app.Topology
        index_to_new_residues : dict
            key : int (index, zero-indexed in chain)
            value : str (three letter residue name)
        Returns
        -------
        mutations : list(str)
            XXX-##-XXX
            three letter WT residue name - id - three letter MUT residue name
            id is based on the protein sequence NOT zero-indexed
        """
        return [r.name+'-'+str(r.id)+'-'+index_to_new_residues[r.index] for r in topology.residues() if r.index in index_to_new_residues]

    def _compute_mutant_key(self, topology, chain_id):
        mutant_key = ''
        chain = ''
        wildtype = self._wildtype
        anychains = []
        anywt_chains = []
        for anychain in topology.chains():
            anychains.append(anychain.id)
            if anychain.id == chain_id:
                chain = anychain
                break
        if not chain:
            raise Exception("Chain %s not found. Available chains: %s" %(chain_id, anychains))
        for anywt_chain in wildtype.chains():
            anywt_chains.append(anywt_chain)
            if anywt_chain.id == chain_id:
                wt_chain = anywt_chain
                break
        if not wt_chain:
            raise Exception("Chain %s not found. Available chains: %s" %(chain_id, anywt_chains))
        for wt_res, res in zip(wt_chain.residues(), chain.residues()):
            if wt_res.name != res.name:
                if mutant_key:
                    mutant_key+='-'
                mutant_key += str(wt_res.name)+str(res.id)+str(res.name)
        if not mutant_key:
            mutant_key = 'WT'
        return mutant_key

class PeptideLibraryEngine(PolymerProposalEngine):
    """
    Note: The PeptideLibraryEngine currently doesn't work because PolymerProposalEngine has been modified to handle
    only one mutation at a time (in accordance with the geometry engine).
    Arguments
    --------
    system_generator : SystemGenerator
    library : list of strings
        each string is a 1-letter-code list of amino acid sequence
    chain_id : str
        id of the chain to mutate
        (using the first chain with the id, if there are multiple)
    proposal_metadata : dict -- OPTIONAL
        Contains information necessary to initialize proposal engine
    """

    def __init__(self, system_generator, library, chain_id, proposal_metadata=None, verbose=False, always_change=True):
        super(PeptideLibraryEngine,self).__init__(system_generator, chain_id, proposal_metadata=proposal_metadata, verbose=verbose, always_change=always_change)
        self._library = library
        self._ff = system_generator.forcefield
        self._templates = self._ff._templates

    def _choose_mutant(self, topology, metadata):
        """
        Used when library of pepide sequences has been provided
        Assume (for now) uniform probability of selecting each peptide
        Arguments
        ---------
        topology : simtk.openmm.app.Topology
        chain_id : str
        allowed_mutations : list(list(tuple))
            list of allowed mutant states; each entry in the list is a list because multiple mutations may be desired
            tuple : (str, str) -- residue id and three-letter amino acid code of desired mutant
        Returns
        -------
        index_to_new_residues : dict
            key : int (index, zero-indexed in chain)
            value : str (three letter residue name)
        metadata : dict
            has not been altered
        """
        library = self._library

        index_to_new_residues = dict()

        # chain : simtk.openmm.app.topology.Chain
        chain_id = self._chain_id
        chain_found = False
        for chain in topology.chains():
            if chain.id == chain_id:
                chain_found = True
                break
        if not chain_found:
            chains = [ chain.id for chain in topology.chains() ]
            raise Exception("Chain '%s' not found in Topology. Chains present are: %s" % (chain_id, str(chains)))
        # location_prob : np.array, probability value for each residue location (uniform)
        location_prob = np.array([1.0/len(library) for i in range(len(library))])
        proposed_location = np.random.choice(range(len(library)), p=location_prob)
        for residue_index, residue_one_letter in enumerate(library[proposed_location]):
            # original_residue : simtk.openmm.app.topology.Residue
            original_residue = chain._residues[residue_index]
            residue_name = self._one_to_three_letter_code(residue_one_letter)
            if original_residue.name == residue_name:
                continue
            # index_to_new_residues : dict, key : int (index of residue, 0-indexed), value : str (three letter residue name)
            index_to_new_residues[residue_index] = residue_name
            if residue_name == 'HIS':
                his_state = ['HIE','HID']
                his_prob = np.array([0.5 for i in range(len(his_state))])
                his_choice = np.random.choice(range(len(his_state)),p=his_prob)
                index_to_new_residues[residue_index] = his_state[his_choice]

        # index_to_new_residues : dict, key : int (index of residue, 0-indexed), value : str (three letter residue name)
        return index_to_new_residues, metadata

    def _one_to_three_letter_code(self, residue_one_letter):
        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'
        }
        return three_letter_code[residue_one_letter]

class SystemGenerator(object):
    """
    This is a utility class to generate OpenMM Systems from
    topology objects.

    Parameters
    ----------
    forcefields_to_use : list of string
        List of the names of ffxml files that will be used in system creation.
    forcefield_kwargs : dict of arguments to createSystem, optional
        Allows specification of various aspects of system creation.
    metadata : dict, optional
        Metadata associated with the SystemGenerator.
    use_antechamber : bool, optional, default=True
        If True, will add the GAFF residue template generator.
    barostat : MonteCarloBarostat, optional, default=None
        If provided, a matching barostat will be added to the generated system.
    particle_charges : bool, optional, default=True
        If False, particle charges will be zeroed
    exception_charges : bool, optional, default=True
        If False, exception charges will be zeroed.
    particle_epsilon : bool, optional, default=True
        If False, particle LJ epsilon will be zeroed.
    exception_epsilon : bool, optional, default=True
        If False, exception LJ epsilon will be zeroed.
    torsions : bool, optional, default=True
        If False, torsions will be zeroed.
    """

    def __init__(self, forcefields_to_use, forcefield_kwargs=None, metadata=None, use_antechamber=True, barostat=None,
        particle_charge=True, exception_charge=True, particle_epsilon=True, exception_epsilon=True,
        torsions=True, angles=True):
        self._forcefield_xmls = forcefields_to_use
        self._forcefield_kwargs = forcefield_kwargs if forcefield_kwargs is not None else {}
        self._forcefield = app.ForceField(*self._forcefield_xmls)
        if use_antechamber:
            self._forcefield.registerTemplateGenerator(forcefield_generators.gaffTemplateGenerator)
        if 'removeCMMotion' not in self._forcefield_kwargs:
            self._forcefield_kwargs['removeCMMotion'] = False
        self._barostat = None
        if barostat is not None:
            pressure = barostat.getDefaultPressure()
            if hasattr(barostat, 'getDefaultTemperature'):
                temperature = barostat.getDefaultTemperature()
            else:
                temperature = barostat.getTemperature()
            frequency = barostat.getFrequency()
            self._barostat = (pressure, temperature, frequency)

        self._particle_charge = particle_charge
        self._exception_charge = exception_charge
        self._particle_epsilon = particle_epsilon
        self._exception_epsilon = exception_epsilon
        self._torsions = torsions

    def getForceField(self):
        """
        Return the associated ForceField object.

        Returns
        -------
        forcefield : simtk.openmm.app.ForceField
            The current ForceField object.
        """
        return self._forcefield

    def build_system(self, new_topology, check_system=False):
        """
        Build a system from the new_topology, adding templates
        for the molecules in oemol_list

        Parameters
        ----------
        new_topology : simtk.openmm.app.Topology object
            The topology of the system
        check_system : book, optional, default=False`
            If True, will check system for issues following creation

        Returns
        -------
        new_system : openmm.System
            A system object generated from the topology
        """
        # TODO: Write some debug info if exception is raised
        system = self._forcefield.createSystem(new_topology, **self._forcefield_kwargs)

        # Turn off various force classes for debugging if requested
        for force in system.getForces():
            if force.__class__.__name__ == 'NonbondedForce':
                for index in range(force.getNumParticles()):
                    charge, sigma, epsilon = force.getParticleParameters(index)
                    if not self._particle_charge:
                        charge *= 0
                    if not self._particle_epsilon:
                        epsilon *= 0
                    force.setParticleParameters(index, charge, sigma, epsilon)
                for index in range(force.getNumExceptions()):
                    p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(index)
                    if not self._exception_charge:
                        chargeProd *= 0
                    if not self._exception_epsilon:
                        epsilon *= 0
                    force.setExceptionParameters(index, p1, p2, chargeProd, sigma, epsilon)
            elif force.__class__.__name__ == 'PeriodicTorsionForce':
                for index in range(force.getNumTorsions()):
                    p1, p2, p3, p4, periodicity, phase, K = force.getTorsionParameters(index)
                    if not self._torsions:
                        K *= 0
                    force.setTorsionParameters(index, p1, p2, p3, p4, periodicity, phase, K)

        # Add barostat if requested.
        if self._barostat is not None:
            MAXINT = np.iinfo(np.int32).max
            barostat = openmm.MonteCarloBarostat(*self._barostat)
            seed = np.random.randint(MAXINT)
            barostat.setRandomNumberSeed(seed)
            system.addForce(barostat)

        # See if any torsions have duplicate atoms.
        if check_system:
            from perses.tests import utils
            utils.check_system(system)

        return system

    @property
    def ffxmls(self):
        return self._forcefield_xmls

    @property
    def forcefield(self):
        return self._forcefield

class DummyForceField(object):
    """
    Dummy force field that can add basic parameters to any system for testing purposes.
    """
    def createSystem(self, topology, **kwargs):
        """
        Create a System object with simple parameters from the provided Topology

        Any kwargs are ignored.

        Parameters
        ----------
        topology : simtk.openmm.app.Topology
            The Topology to be parameterized

        Returns
        -------
        system : simtk.openmm.System
            The System object
        """
        from openmmtools.constants import kB
        kT = kB * 300*unit.kelvin

        # Create a System
        system = openmm.System()

        # Add particles
        nonbonded = openmm.CustomNonbondedForce('100/(r/0.1)^4')
        nonbonded.setNonbondedMethod(openmm.CustomNonbondedForce.CutoffNonPeriodic);
        nonbonded.setCutoffDistance(1*unit.nanometer)
        system.addForce(nonbonded)
        mass = 12.0 * unit.amu
        for atom in topology.atoms():
            nonbonded.addParticle([])
            system.addParticle(mass)

        # Build a list of which atom indices are bonded to each atom
        bondedToAtom = []
        for atom in topology.atoms():
            bondedToAtom.append(set())
        for (atom1, atom2) in topology.bonds():
            bondedToAtom[atom1.index].add(atom2.index)
            bondedToAtom[atom2.index].add(atom1.index)
        return bondedToAtom

        # Add bonds
        bond_force = openmm.HarmonicBondForce()
        r0 = 1.0 * unit.angstroms
        sigma_r = 0.1 * unit.angstroms
        Kr = kT / sigma_r**2
        for atom1, atom2 in topology.bonds():
            bond_force.addBond(atom1.index, atom2.index, r0, Kr)
        system.addForce(bond_force)

        # Add angles
        uniqueAngles = set()
        for bond in topology.bonds():
            for atom in bondedToAtom[bond.atom1]:
                if atom != bond.atom2:
                    if atom < bond.atom2:
                        uniqueAngles.add((atom, bond.atom1, bond.atom2))
                    else:
                        uniqueAngles.add((bond.atom2, bond.atom1, atom))
            for atom in bondedToAtom[bond.atom2]:
                if atom != bond.atom1:
                    if atom > bond.atom1:
                        uniqueAngles.add((bond.atom1, bond.atom2, atom))
                    else:
                        uniqueAngles.add((atom, bond.atom2, bond.atom1))
        angles = sorted(list(uniqueAngles))
        theta0 = 109.5 * unit.degrees
        sigma_theta = 10 * unit.degrees
        Ktheta = kT / sigma_theta**2
        angle_force = openmm.HarmonicAngleForce()
        for (atom1, atom2, atom3) in angles:
            angles.addAngle(atom1.index, atom2.index, atom3.index, theta0, Ktheta)
        system.addForce(angle_force)

        # Make a list of all unique proper torsions
        uniquePropers = set()
        for angle in angles:
            for atom in bondedToAtom[angle[0]]:
                if atom not in angle:
                    if atom < angle[2]:
                        uniquePropers.add((atom, angle[0], angle[1], angle[2]))
                    else:
                        uniquePropers.add((angle[2], angle[1], angle[0], atom))
            for atom in bondedToAtom[angle[2]]:
                if atom not in angle:
                    if atom > angle[0]:
                        uniquePropers.add((angle[0], angle[1], angle[2], atom))
                    else:
                        uniquePropers.add((atom, angle[2], angle[1], angle[0]))
        propers = sorted(list(uniquePropers))
        torsion_force = openmm.PeriodicTorsionForce()
        periodicity = 3
        phase = 0.0 * unit.degrees
        Kphi = 0.0 * kT
        for (atom1, atom2, atom3, atom4) in propers:
            torsion_force.add_torsion(atom1.index, atom2.index, atom3.index, atom4.index, periodicity, phase, Kphi)
        system.addForce(torsion_force)

        return system

class DummySystemGenerator(SystemGenerator):
    """
    Dummy SystemGenerator that employs a universal simple force field.

    """
    def __init__(self, forcefields_to_use, barostat=None, **kwargs):
        """
        Create a DummySystemGenerator with universal simple force field.

        All parameters except 'barostat' are ignored.

        """
        self._forcefield = DummyForceField()
        self._forcefield_xmls = list()
        self._forcefield_kwargs = dict()
        self._barostat = None
        if barostat is not None:
            pressure = barostat.getDefaultPressure()
            if hasattr(barostat, 'getDefaultTemperature'):
                temperature = barostat.getDefaultTemperature()
            else:
                temperature = barostat.getTemperature()
            frequency = barostat.getFrequency()
            self._barostat = (pressure, temperature, frequency)


class SmallMoleculeSetProposalEngine(AtomMapper, ProposalEngine):
    """
    This class proposes new small molecules from a prespecified set. It uses
    uniform proposal probabilities, but can be extended. The user is responsible
    for providing a list of smiles that can be interconverted! The class includes
    extra functionality to assist with that (it is slow).

    Parameters
    ----------
    list_of_smiles : list of string
        list of smiles that will be sampled
    system_generator : SystemGenerator object
        SystemGenerator initialized with the appropriate forcefields
    residue_name : str
        The name that will be used for small molecule residues in the topology
    proposal_metadata : dict
        metadata for the proposal engine
    storage : NetCDFStorageView, optional, default=None
        If specified, write statistics to this storage
    """

    def __init__(self, list_of_oemols, system_generator, residue_name='MOL', storage=None, **kwargs):
        super(SmallMoleculeSetProposalEngine, self).__init__(list_of_oemols=list_of_oemols,system_generator=system_generator, **kwargs)
        # This needs to be exposed, and only set in one place

        self._n_molecules = len(self.list_of_oemols)

        self._residue_name = residue_name
        self._generated_systems = dict()
        self._generated_topologies = dict()
        self._matches = dict()

        self._list_of_smiles = []
        from perses.utils.openeye import createSMILESfromOEMol
        for mol in self.list_of_oemols:
            smiles = createSMILESfromOEMol(mol)
            self._list_of_smiles.append(SmallMoleculeSetProposalEngine.canonicalize_smiles(smiles))

        self._storage = storage
        if self._storage is not None:
            self._storage = NetCDFStorageView(storage, modname=self.__class__.__name__)

        _logger.info(f"creating probability matrix...")
        self._probability_matrix = self._calculate_probability_matrix()


    def propose(self, current_system, current_topology,
                current_mol_id=0, proposed_mol_id=None, current_metadata=None):
        """
        Propose the next state, given the current state

        Parameters
        ----------
        current_system : openmm.System object
            the system of the current state
        current_topology : app.Topology object
            the topology of the current state
        current_metadata : dict
            dict containing current smiles as a key
        current_mol_id : int, optional, default=0
            Index of starting oemol, default is first in list
        proposed_mol_id : int, optional, default=None
            If specified, index of oemol to propose, if None, an oemol from the list is chosen

        Returns
        -------
        proposal : TopologyProposal object
           topology proposal object
        """
        self.current_mol_id = current_mol_id
        if len(self.list_of_oemols) == 2:
            # only two molecules so...
            if self.current_mol_id == 0:
                self.proposed_mol_id = 1
            elif self.current_mol_id == 1:
                self.proposed_mol_id = 0
        else:
            self.proposed_mol_id = proposed_mol_id
        self.current_molecule = self.list_of_oemols[self.current_mol_id]

        # Remove the small molecule from the current Topology object
        _logger.info(f"creating current receptor topology by removing small molecule from current topology...")
        current_receptor_topology = self._remove_small_molecule(current_topology)

        # Find the initial atom index of the small molecule in the current topology
        old_mol_start_index, len_old_mol = self._find_mol_start_index(current_topology)
        self.old_mol_start_index = old_mol_start_index
        self.len_old_mol = len_old_mol
        _logger.info(f"small molecule start index: {old_mol_start_index}")
        _logger.info(f"small molecule has {len_old_mol} atoms.")

        # Determine atom indices of the small molecule in the current topology
        old_alchemical_atoms = range(old_mol_start_index, len_old_mol)
        _logger.info(f"old alchemical atom indices: {old_alchemical_atoms}")

        # Select the next molecule SMILES given proposal probabilities
        if self.proposed_mol_id is None:
            _logger.info(f"the proposed oemol is not specified; proposing a new molecule from proposal matrix P(M_new | M_old)...")
            self.proposed_mol_id, self.proposed_molecule, logp_proposal = self._propose_molecule(current_system, current_topology, self.current_mol_id)
        else:
            self.proposed_molecule = self.list_of_oemols[self.proposed_mol_id]
            proposed_mol_smiles = self._list_of_smiles[self.proposed_mol_id]
            _logger.info(f"proposed mol detected with smiles {proposed_mol_smiles} and logp_proposal of 0.0")
            logp_proposal = 0.0

        _logger.info(f"conducting proposal from {self._list_of_smiles[self.current_mol_id]} to {self._list_of_smiles[self.proposed_mol_id]}...")

        # Build the new Topology object, including the proposed molecule
        _logger.info(f"building new topology with proposed molecule and current receptor topology...")
        new_topology = self._build_new_topology(current_receptor_topology, self.proposed_molecule)
        new_mol_start_index, len_new_mol = self._find_mol_start_index(new_topology)
        self.new_mol_start_index = new_mol_start_index
        self.len_new_mol = len_new_mol
        _logger.info(f"new molecule has a start index of {new_mol_start_index} and {len_new_mol} atoms.")

        # Generate an OpenMM System from the proposed Topology
        _logger.info(f"proceeding to build the new system from the new topology...")
        # TODO: Remove build_system() branch once we convert entirely to new openmm-forcefields SystemBuilder
        if hasattr(self._system_generator, 'create_system'):
            new_system = self._system_generator.create_system(new_topology)
        else:
            new_system = self._system_generator.build_system(new_topology)

        # Determine atom mapping between old and new molecules
        _logger.info(f"determining atom map between old and new molecules...")
        if not self.atom_map:
            _logger.info(f"the atom map is not specified; proceeding to generate an atom map...")
            mol_atom_map = self._get_mol_atom_map()
        else:
            _logger.info(f"atom map is pre-determined as {self.atom_map}")
            mol_atom_map = self.atom_map

        # Adjust atom mapping indices for the presence of the receptor
        adjusted_atom_map = {}
        for (key, value) in mol_atom_map.items():
            adjusted_atom_map[key+new_mol_start_index] = value + old_mol_start_index

        # Incorporate atom mapping of all environment atoms
        old_mol_offset = len_old_mol
        for i in range(new_mol_start_index):
            if i >= old_mol_start_index:
                old_idx = i + old_mol_offset
            else:
                old_idx = i
            adjusted_atom_map[i] = old_idx

        # now to correct for possible constraint problems
        adjusted_atom_map = SmallMoleculeSetProposalEngine._constraint_repairs(adjusted_atom_map, current_system, new_system, current_topology, new_topology)
        non_offset_new_to_old_atom_map = copy.deepcopy(adjusted_atom_map)
        # TODO is the following line needed? It doesn't seem to be used
        min_keys, min_values = min(non_offset_new_to_old_atom_map.keys()), min(non_offset_new_to_old_atom_map.values())
        self.non_offset_new_to_old_atom_map = mol_atom_map

        # Create the TopologyProposal onbject
        proposal = TopologyProposal(logp_proposal=logp_proposal,
                                    new_to_old_atom_map=adjusted_atom_map,
                                    old_topology=current_topology,
                                    new_topology=new_topology,
                                    old_system=current_system,
                                    new_system=new_system,
                                    old_alchemical_atoms=old_alchemical_atoms,
                                    old_chemical_state_key=self._list_of_smiles[self.current_mol_id],
                                    new_chemical_state_key=self._list_of_smiles[self.proposed_mol_id],
                                    old_residue_name=self._residue_name,
                                    new_residue_name=self._residue_name)

        ndelete = proposal.old_system.getNumParticles() - len(proposal.old_to_new_atom_map.keys())
        ncreate = proposal.new_system.getNumParticles() - len(proposal.old_to_new_atom_map.keys())
        _logger.info(f'Proposed transformation would delete {ndelete} atoms and create {ncreate} atoms.')

        return proposal

    @staticmethod
    def canonicalize_smiles(smiles):
        """
        Convert a SMILES string into canonical isomeric smiles

        Parameters
        ----------
        smiles : str
            Any valid SMILES for a molecule

        Returns
        -------
        iso_can_smiles : str
            OpenEye isomeric canonical smiles corresponding to the input
        """
        # TODO can this go in perses/utils/openeye?
        mol = oechem.OEMol()
        oechem.OESmilesToMol(mol, smiles)
        oechem.OEAddExplicitHydrogens(mol)
        iso_can_smiles = oechem.OECreateSmiString(mol, OESMILES_OPTIONS)
        return iso_can_smiles

    def _topology_to_smiles(self, topology):
        """
        Get the smiles string corresponding to a specific residue in an
        OpenMM Topology

        Parameters
        ----------
        topology : app.Topology
            The topology containing the molecule of interest

        Returns
        -------
        smiles_string : string
            an isomeric canonicalized SMILES string representing the molecule
        oemol : oechem.OEMol object
            molecule
        """
        # TODO can this go in perses/utils/openeye?
        molecule_name = self._residue_name
        _logger.info(f"\tmolecule name specified from residue: {self._residue_name}.")

        matching_molecules = [res for res in topology.residues() if res.name[:3] == molecule_name[:3]]  # Find residue in topology by searching for residues with name "MOL"
        if len(matching_molecules) > 1:
            raise ValueError("More than one residue with the same name!")
        if len(matching_molecules) == 0:
            raise ValueError(f"No residue found with the resname {molecule_name[:3]}")
        mol_res = matching_molecules[0]
        oemol = forcefield_generators.generateOEMolFromTopologyResidue(mol_res)
        _logger.info(f"\toemol generated!")
        smiles_string = oechem.OECreateSmiString(oemol, OESMILES_OPTIONS)
        final_smiles_string = smiles_string
        _logger.info(f"\tsmiles generated from oemol: {final_smiles_string}")
        return final_smiles_string, oemol

    def compute_state_key(self, topology):
        """
        Given a topology, come up with a state key string.
        For this class, the state key is an isomeric canonical SMILES.

        Parameters
        ----------
        topology : app.Topology object
            The topology object in question.

        Returns
        -------
        chemical_state_key : str
            isomeric canonical SMILES

        """
        chemical_state_key, _ = self._topology_to_smiles(topology)
        return chemical_state_key

    def _find_mol_start_index(self, topology):
        """
        Find the starting index of the molecule in the topology.
        Throws an exception if resname is not present.

        Parameters
        ----------
        topology : app.Topology object
            The topology containing the molecule

        Returns
        -------
        mol_start_idx : int
            start index of the molecule
        mol_length : int
            the number of atoms in the molecule
        """
        resname = self._residue_name
        mol_residues = [res for res in topology.residues() if res.name[:3]==resname[:3]]  # Find the residue by searching for residues with "MOL"
        if len(mol_residues)!=1:
            raise ValueError("There must be exactly one residue with a specific name in the topology. Found %d residues with name '%s'" % (len(mol_residues), resname))
        mol_residue = mol_residues[0]
        atoms = list(mol_residue.atoms())
        mol_start_idx = atoms[0].index
        return mol_start_idx, len(list(atoms))

    def _build_new_topology(self, current_receptor_topology, oemol_proposed):
        """
        Construct a new topology
        Parameters
        ----------
        oemol_proposed : oechem.OEMol object
            the proposed OEMol object
        current_receptor_topology : app.Topology object
            The current topology without the small molecule

        Returns
        -------
        new_topology : app.Topology object
            A topology with the receptor and the proposed oemol
        mol_start_index : int
            The first index of the small molecule
        """
        _logger.info(f"\tsetting proposed oemol title to {self._residue_name}")
        oemol_proposed.SetTitle(self._residue_name)
        _logger.info(f"\tcreating mol topology from oemol...")
        mol_topology = forcefield_generators.generateTopologyFromOEMol(oemol_proposed)
        new_topology = app.Topology()
        _logger.info(f"\tappending current receptor topology to new mol topology...")
        append_topology(new_topology, current_receptor_topology)
        append_topology(new_topology, mol_topology)
        # Copy periodic box vectors.
        if current_receptor_topology._periodicBoxVectors != None:
            _logger.info(f"\tperiodic box vectors of the current receptor is specified; copying to new topology...")
            new_topology._periodicBoxVectors = copy.deepcopy(current_receptor_topology._periodicBoxVectors)

        return new_topology

    def _remove_small_molecule(self, topology):
        """
        This method removes the small molecule parts of a topology from
        a protein+small molecule complex based on the name of the
        small molecule residue

        Parameters
        ----------
        topology : app.Topology
            Topology with the appropriate residue name.

        Returns
        -------
        receptor_topology : app.Topology
            Topology without small molecule
        """
        receptor_topology = app.Topology()
        append_topology(receptor_topology, topology, exclude_residue_name=self._residue_name)
        # Copy periodic box vectors.
        if topology._periodicBoxVectors != None:
            receptor_topology._periodicBoxVectors = copy.deepcopy(topology._periodicBoxVectors)
        return receptor_topology

    @staticmethod
    def _constraint_repairs(atom_map, old_system, new_system, old_topology, new_topology):
        """
        Given an adjusted atom map (corresponding to the true indices of the new: old atoms in their respective systems), iterate through all of the
        atoms in the map that are hydrogen and check if the constraint length changes; if so, we do not map.
        """
        old_hydrogens = list(atom.index for atom in old_topology.atoms() if atom.element == app.Element.getByAtomicNumber(1))
        new_hydrogens = list(atom.index for atom in new_topology.atoms() if atom.element == app.Element.getByAtomicNumber(1))
        #wrapping constraints
        old_constraints, new_constraints = {}, {}
        for idx in range(old_system.getNumConstraints()):
            atom1, atom2, length = old_system.getConstraintParameters(idx)
            if atom1 in old_hydrogens:
                old_constraints[atom1] = length
            elif atom2 in old_hydrogens:
                old_constraints[atom2] = length

        for idx in range(new_system.getNumConstraints()):
            atom1, atom2, length = new_system.getConstraintParameters(idx)
            if atom1 in new_hydrogens:
                new_constraints[atom1] = length
            elif atom2 in new_hydrogens:
                new_constraints[atom2] = length

         #iterate through the atom indices in the new_to_old map, check bonds for pairs, and remove appropriate matches
        to_delete = []
        for new_index, old_index in atom_map.items():
            if new_index in new_constraints.keys() and old_index in old_constraints.keys(): # both atom indices are hydrogens
                old_length, new_length = old_constraints[old_index], new_constraints[new_index]
                if not old_length == new_length: #then we have to remove it from
                    to_delete.append(new_index)

        for idx in to_delete:
            del atom_map[idx]

        return atom_map

    def _propose_molecule(self, system, topology, current_mol_id, exclude_self=False):
        """
        Propose a new molecule given the current molecule.

        The current scheme uses a probability matrix computed via _calculate_probability_matrix.

        Arguments
        ---------
        system : simtk.openmm.System object
            The current system
        topology : simtk.openmm.app.Topology object
            The current topology
        positions : [n, 3] np.ndarray of floats (Quantity nm)
            The current positions of the system
        current_mol_id : int
            The index of the current molecule
        exclude_self : bool, optional, default=True
            If True, exclude self-transitions

        Returns
        -------
        proposed_mol_id : int
             The index of the proposed molecule
        mol : oechem.OEMol
            The next molecule to simulate
        logp_proposal : float
            contribution from the chemical proposal to the log probability of acceptance (Eq. 36 for hybrid; Eq. 53 for two-stage)
            log [P(Mold | Mnew) / P(Mnew | Mold)]
        """
        # Compute contribution from the chemical proposal to the log probability of acceptance (Eq. 36 for hybrid; Eq. 53 for two-stage)
        # log [P(Mold | Mnew) / P(Mnew | Mold)]

        # Retrieve the current molecule index

        # Propose a new molecule
        molecule_probabilities = self._probability_matrix[current_mol_id, :]
        _logger.info(f"\tmolecule probabilities: {molecule_probabilities}")
        proposed_mol_id = np.random.choice(range(len(self._list_of_smiles)), p=molecule_probabilities)
        _logger.info(f"\tproposed molecule index: {proposed_mol_id}")
        reverse_probability = self._probability_matrix[proposed_mol_id, current_mol_id]
        forward_probability = molecule_probabilities[proposed_mol_id]
        _logger.info(f"\tforward probability: {forward_probability}")
        _logger.info(f"\treverse probability: {reverse_probability}")
        proposed_smiles = self._list_of_smiles[proposed_mol_id]
        _logger.info(f"\tproposed molecule smiles: {proposed_smiles}")
        proposed_mol = self.list_of_oemols[proposed_mol_id]
        logp = np.log(reverse_probability) - np.log(forward_probability)
        _logger.info(f"\tlogP proposal: {logp}")
        return proposed_mol_id, proposed_mol, logp

    def _calculate_probability_matrix(self):
        """
        Calculate the matrix of probabilities of choosing A | B
        based on normalized MCSS overlap. Does not check for torsions!
        Parameters
        ----------
        oemol_list : list of oechem.OEMol
            list of oemols to be potentially selected

        Returns
        -------
        probability_matrix : [n, n] np.ndarray
            probability_matrix[Mold, Mnew] is the probability of choosing molecule Mnew given the current molecule is Mold

        """
        n_mols = len(self.list_of_oemols)
        probability_matrix = np.zeros([n_mols, n_mols])
        for i in range(n_mols):
            for j in range(i):
                atom_map = self._get_mol_atom_map(indexA=i, indexB=j)
                if not atom_map:
                    n_atoms_matching = 0
                    continue
                n_atoms_matching = len(atom_map.keys())
                probability_matrix[i, j] = n_atoms_matching
                probability_matrix[j, i] = n_atoms_matching
        #normalize the rows:
        for i in range(n_mols):
            row_sum = np.sum(probability_matrix[i, :])
            try:
                probability_matrix[i, :] /= row_sum
            except ZeroDivisionError:
                print("One molecule is completely disconnected!")
                raise

        if self._storage:
            self._storage.write_object('molecule_smiles_list', self._list_of_smiles)
            self._storage.write_array('probability_matrix', probability_matrix)

        return probability_matrix

    @property
    def chemical_state_list(self):
        return self._list_of_smiles

    @staticmethod
    def clean_molecule_list(smiles_list, atom_opts, bond_opts):
        """
        A utility function to determine which molecules can be proposed
        from any other molecule.

        Parameters
        ----------
        smiles_list
        atom_opts
        bond_opts

        Returns
        -------
        safe_smiles
        removed_smiles
        """
        from perses.tests.utils import createSystemFromSMILES
        import itertools
        from perses.rjmc.geometry import ProposalOrderTools
        from perses.tests.test_geometry_engine import oemol_to_openmm_system
        safe_smiles = set()
        smiles_pairs = set()
        smiles_set = set(smiles_list)

        for mol1, mol2 in itertools.combinations(smiles_list, 2):
            smiles_pairs.add((mol1, mol2))

        for smiles_pair in smiles_pairs:
            mol1, sys1, pos1, top1 = createSystemFromSMILES(smiles_pair[0])
            mol2, sys2, pos2, top2 = createSystemFromSMILES(smiles_pair[1])

            new_to_old_atom_map = SmallMoleculeSetProposalEngine._get_mol_atom_map()
            if not new_to_old_atom_map:
                continue
            top_proposal = TopologyProposal(new_topology=top2, old_topology=top1, new_system=sys2, old_system=sys1, new_to_old_atom_map=new_to_old_atom_map, new_chemical_state_key='e', old_chemical_state_key='w')
            proposal_order = ProposalOrderTools(top_proposal)
            try:
                forward_order = proposal_order.determine_proposal_order(direction='forward')
                reverse_order = proposal_order.determine_proposal_order(direction='reverse')
                safe_smiles.add(smiles_pair[0])
                safe_smiles.add(smiles_pair[1])
                _logger.info("Adding %s and %s" % (smiles_pair[0], smiles_pair[1]))
            except NoTorsionError:
                pass
        removed_smiles = smiles_set.difference(safe_smiles)
        return safe_smiles, removed_smiles

class PremappedSmallMoleculeSetProposalEngine(SmallMoleculeSetProposalEngine):
    """
    This proposal engine uses the AtomMapper to have all atoms premapped and the proposal distribution pre-formed and checked.
    It is intended to be substantially faster, as well as more robust (having excluded mappings that would not lead to a valid geometry proposal)
    """

    def __init__(self, atom_mapper: AtomMapper, system_generator: SystemGenerator, residue_name: str="MOL", storage: NetCDFStorageView=None):
        self._atom_mapper = atom_mapper
        self._atom_mapper.map_all_molecules()
        self._atom_mapper.generate_and_check_proposal_matrix()
        self._proposal_matrix = self._atom_mapper.proposal_matrix

        self._n_molecules = self._atom_mapper.n_molecules

        super(PremappedSmallMoleculeSetProposalEngine, self).__init__(self._atom_mapper.smiles_list, system_generator, residue_name=residue_name,
                 proposal_metadata=None, storage=storage,
                 always_change=True)

    def propose(self, current_system, current_topology, current_smiles=None, proposed_mol=None, map_index=None, current_metadata=None):
        """
        Propose the next state, given the current state

        Parameters
        ----------
        current_system : openmm.System object
            the system of the current state
        current_topology : app.Topology object
            the topology of the current state
        current_smiles : str, default None
            Specify the current SMILES string to avoid perceiving it from the topology. If None, perceive from topology.
        proposed_mol : oechem.OEMol, optional
            the molecule to propose. If None, choose randomly based on the current molecule
        map_index : int, default None
            The index of the atom map to use. If None, choose randomly. Otherwise, use map idx of map_index mod n_maps
        current_metadata : dict
            dict containing current smiles as a key

        Returns
        -------
        proposal : TopologyProposal object
           topology proposal object
        """
        # Determine SMILES string for current small molecule if the SMILES isn't specified
        if current_smiles is None:
            current_mol_smiles, _ = self._topology_to_smiles(current_topology)
        else:
            current_mol_smiles = current_smiles

        # Remove the small molecule from the current Topology object
        current_receptor_topology = self._remove_small_molecule(current_topology)

        # Find the initial atom index of the small molecule in the current topology
        old_mol_start_index, len_old_mol = self._find_mol_start_index(current_topology)

        # Determine atom indices of the small molecule in the current topology
        old_alchemical_atoms = range(old_mol_start_index, len_old_mol)

        # Select the next molecule SMILES given proposal probabilities
        current_mol_index = self._atom_mapper.get_smiles_index(current_mol_smiles)

        #If we aren't specifying a proposed molecule, then randomly propose one:
        if proposed_mol is None:
            #get probability vector for proposal
            proposal_probability = self._proposal_matrix[current_mol_index, :]

            #propose next index
            proposed_index = np.random.choice(range(self._n_molecules), p=proposal_probability)

            #proposal logp
            proposed_logp = np.log(proposal_probability[proposed_index])

            #reverse proposal logp
            reverse_logp = np.log(self._proposal_matrix[proposed_index, current_mol_index])

            #logp overall of proposal
            logp_proposal = reverse_logp - proposed_logp

            #get the oemol corresponding to the proposed molecule:
            proposed_mol_smiles = self._atom_mapper.smiles_list[proposed_index]
            proposed_mol = self._atom_mapper.get_oemol_from_smiles(proposed_mol_smiles)

        else:
            logp_proposal = 0.0
            proposed_mol_smiles = oechem.OECreateSmiString(proposed_mol, OESMILES_OPTIONS)

        #You will get a weird error if you don't assign atom names.
        oechem.OETriposAtomNames(proposed_mol)

        # Build the new Topology object, including the proposed molecule
        new_topology = self._build_new_topology(current_receptor_topology, proposed_mol)
        new_mol_start_index, len_new_mol = self._find_mol_start_index(new_topology)

        # Generate an OpenMM System from the proposed Topology
        # TODO: Remove build_system() branch once we convert entirely to new openmm-forcefields SystemBuilder
        if hasattr(self._system_generator, 'create_system'):
            new_system = self._system_generator.create_system(new_topology)
        else:
            new_system = self._system_generator.build_system(new_topology)

        # Determine atom mapping between old and new molecules
        mol_atom_maps = self._atom_mapper.get_atom_maps(current_mol_smiles, proposed_mol_smiles)

        #If no map index is given, just randomly choose one:
        if map_index is None:
            mol_atom_map = np.random.choice(mol_atom_maps)

        #Otherwise, pick the map whose index is the specified index mod n_maps
        else:
            n_atom_maps = len(mol_atom_maps)
            index_to_choose = map_index % n_atom_maps
            mol_atom_map = mol_atom_maps[index_to_choose]

        # Adjust atom mapping indices for the presence of the receptor
        adjusted_atom_map = {}
        for (key, value) in mol_atom_map.items():
            adjusted_atom_map[key+new_mol_start_index] = value + old_mol_start_index

        # Incorporate atom mapping of all environment atoms
        old_mol_offset = len_old_mol
        for i in range(new_mol_start_index):
            if i >= old_mol_start_index:
                old_idx = i + old_mol_offset
            else:
                old_idx = i
            adjusted_atom_map[i] = old_idx

        # Create the TopologyProposal onbject
        proposal = TopologyProposal(logp_proposal=logp_proposal, new_to_old_atom_map=adjusted_atom_map,
            old_topology=current_topology, new_topology=new_topology,
            old_system=current_system, new_system=new_system,
            old_alchemical_atoms=old_alchemical_atoms,
            old_chemical_state_key=current_mol_smiles, new_chemical_state_key=proposed_mol_smiles)

        ndelete = proposal.old_system.getNumParticles() - len(proposal.old_to_new_atom_map.keys())
        ncreate = proposal.new_system.getNumParticles() - len(proposal.old_to_new_atom_map.keys())
        _logger.info('Proposed transformation would delete %d atoms and create %d atoms.' % (ndelete, ncreate))

        return proposal


class TwoMoleculeSetProposalEngine(SmallMoleculeSetProposalEngine):
    """
    Dummy proposal engine that always chooses the new molecule of the two, but uses the base class's atom mapping
    functionality.
    """

    def __init__(self, list_of_oemols, system_generator, residue_name='MOL', atom_expr=None, bond_expr=None, proposal_metadata=None, storage=None, always_change=True, atom_map=None, **kwargs):
        if len(list_of_oemols) != 2:
            _logger.warning(f'TwoMoleculeSetProposalEngine only takes TWO molecules as input')
            _logger.warning(f'Only first two oemols of the list provided (of {len(list_of_oemols)}) will be considered')
        self.current_molecule = list_of_oemols[0]
        self.proposed_molecule = list_of_oemols[1]
        self._old_mol_smiles = oechem.OECreateSmiString(self.current_molecule, OESMILES_OPTIONS)
        self._new_mol_smiles = oechem.OECreateSmiString(self.proposed_molecule, OESMILES_OPTIONS)
        self.current_mol_id = 0
        self.proposed_mol_id = 1

        super(TwoMoleculeSetProposalEngine, self).__init__(list_of_oemols=[self.current_molecule, self.proposed_molecule],system_generator=system_generator, **kwargs)

        self._allow_ring_breaking = False # don't allow ring breaking

    def _propose_molecule(self, system, topology, molecule_smiles, exclude_self=False):
        return self._new_mol_smiles, self.proposed_molecule, 0.0

class NullProposalEngine(SmallMoleculeSetProposalEngine):
    """
    Base class for NaphthaleneProposalEngine and ButantProposalEngine
    Not intended for use on its own

    """
    def __init__(self, system_generator, residue_name="MOL", atom_expr=None, bond_expr=None, proposal_metadata=None, storage=None, always_change=True):
        super(NullProposalEngine, self).__init__(list(), system_generator, residue_name=residue_name)
        self._fake_states = ["A","B"]
        self.smiles = ''

    def propose(self, current_system, current_topology, current_metadata=None):
        """
        Custom proposal for NaphthaleneTestSystem will switch from current naphthalene
        "state" (either 'naphthalene-A' or 'naphthalene-B') to the other

        This proposal engine can only be used with input topology of
        naphthalene, and propose() will first confirm naphthalene is residue "MOL"
        The topology may have water but should not have any other
        carbon-containing residue.

        The "new" system and topology are deep copies of the old, but the atom_map
        is custom defined to only match one of the two rings.

        Arguments:
        ----------
        current_system : openmm.System object
            the system of the current state
        current_topology : app.Topology object
            the topology of the current state
        current_metadata : dict, OPTIONAL
            Not implemented

        Returns:
        -------
        proposal : TopologyProposal object
           topology proposal object

        """
        given_smiles = super(NullProposalEngine, self).compute_state_key(current_topology)
        if self.smiles != given_smiles:
            raise(Exception("{0} can only be used with {1} topology; smiles of given topology is: {2}".format(type(self), self.smiles, given_smiles)))

        old_key = current_topology._state_key
        new_key = [key for key in self._fake_states if key != old_key][0]

        new_topology = app.Topology()
        append_topology(new_topology, current_topology)
        new_topology._state_key = new_key
        new_system = copy.deepcopy(current_system)
        atom_map = self._make_skewed_atom_map(current_topology)
        proposal = TopologyProposal(new_topology=new_topology, new_system=new_system, old_topology=current_topology, old_system=current_system, logp_proposal=0.0,
                                                 new_to_old_atom_map=atom_map, old_chemical_state_key=old_key, new_chemical_state_key=new_key)
        return proposal

    def _make_skewed_atom_map(topology):
        return dict()

    def compute_state_key(self, topology):
        """
        For this test system, the topologies for the two states are
        identical; therefore a custom attribute `_state_key` has
        been added to the topology itself, to track whether a switch
        has been accepted

        compute_state_key will return topology._state_key rather than
        SMILES, because SMILES are identical in the two states.
        """
        return topology._state_key


class NaphthaleneProposalEngine(NullProposalEngine):
    """
    Custom ProposalEngine to use with NaphthaleneTestSystem defines two "states"
    of naphthalene, identified 'naphthalene-A' and 'naphthalene-B', which are
    tracked by adding a custom _state_key attribute to the topology

    Generates TopologyProposal from naphthalene to naphthalene, only matching
    one ring such that geometry must rebuild the other

    Can only be used with input topology of naphthalene

    Constructor Arguments:
        system_generator, SystemGenerator object
            SystemGenerator initialized with the appropriate forcefields
        residue_name, OPTIONAL,  str
            Default = "MOL"
            The name that will be used for small molecule residues in the topology
        atom_expr, OPTIONAL, oechem.OEExprOpts
            Default is None
            Currently not implemented -- would dictate how match is defined
        bond_expr, OPTIONAL, oechem.OEExprOpts
            Default is None
            Currently not implemented -- would dictate how match is defined
        proposal_metadata, OPTIONAL, dict
            Default is None
            metadata for the proposal engine
        storage, OPTIONAL, NetCDFStorageView
            Default is None
            If specified, write statistics to this storage layer
        always_change, OPTIONAL, bool
            Default is True
            Currently not implemented -- will always behave as True
            The proposal will always be from the current "state" to the other
            Self proposals will never be made
    """
    def __init__(self, system_generator, residue_name="MOL", atom_expr=None, bond_expr=None, proposal_metadata=None, storage=None, always_change=True):
        super(NaphthaleneProposalEngine, self).__init__(system_generator, residue_name=residue_name, atom_expr=atom_expr, bond_expr=bond_expr, proposal_metadata=proposal_metadata, storage=storage, always_change=always_change)
        self._fake_states = ["naphthalene-A", "naphthalene-B"]
        self.smiles = 'c1ccc2ccccc2c1'

    def _make_skewed_atom_map(self, topology):
        """
        Custom definition for the atom map between naphthalene and naphthalene

        If a regular atom map was constructed (via oechem.OEMCSSearch), all
        atoms would be matched, and the geometry engine would have nothing
        to add.  This method manually finds one of the two naphthalene rings
        and matches it to itself, rotated 180degrees.  The second ring and
        all hydrogens will have to be repositioned by the geometry engine.

        The rings are identified by finding how many other carbons each
        carbon in the topology is bonded to.  This will be 2 for all carbons
        in naphthalene except the two carbon atoms shared by both rings.
        Starting with these 2 shared atoms, a ring is traced by following
        bonds between adjacent carbons until returning to the shared carbons.

        Arguments:
        ----------
        topology : app.Topology object
            topology of naphthalene
            Only one topology is needed, because current and proposed are
            identical
        Returns:
        --------
        atom_map : dict
            maps the atom indices of carbons from one ring back to themselves
        """
        # make dict of carbons in topology to list of carbons they're bonded to
        carbon_bonds = dict()
        for atom in topology.atoms():
            if atom.element == app.element.carbon:
                carbon_bonds[atom] = set()
        for bond in topology.bonds():
            if bond[0].element == app.element.carbon and bond[1].element == app.element.carbon:
                carbon_bonds[bond[0]].add(bond[1])
                carbon_bonds[bond[1]].add(bond[0])

        # identify the rings by finding 2 middle carbons then tracing their bonds
        find_ring = list()
        for carbon, bounds in carbon_bonds.items():
            if len(bounds) == 3:
                find_ring.append(carbon)
        assert len(find_ring) == 2

        add_next = [carbon for carbon in carbon_bonds[find_ring[1]] if carbon not in find_ring]
        while len(add_next) > 0:
            find_ring.append(add_next[0])
            add_next = [carbon for carbon in carbon_bonds[add_next[0]] if carbon not in find_ring]
        assert len(find_ring) == 6

        # map the ring flipped 180
        atom_map = {find_ring[i].index : find_ring[(i+3)%6].index for i in range(6)}
        return atom_map

class ButaneProposalEngine(NullProposalEngine):
    """
    Custom ProposalEngine to use with ButaneTestSystem defines two "states"
    of butane, identified 'butane-A' and 'butane-B', which are
    tracked by adding a custom _state_key attribute to the topology

    Generates TopologyProposal from butane to butane, only matching
    two carbons such that geometry must rebuild the others

    Can only be used with input topology of butane

    Constructor Arguments:
        system_generator, SystemGenerator object
            SystemGenerator initialized with the appropriate forcefields
        residue_name, OPTIONAL,  str
            Default = "MOL"
            The name that will be used for small molecule residues in the topology
        atom_expr, OPTIONAL, oechem.OEExprOpts
            Default is None
            Currently not implemented -- would dictate how match is defined
        bond_expr, OPTIONAL, oechem.OEExprOpts
            Default is None
            Currently not implemented -- would dictate how match is defined
        proposal_metadata, OPTIONAL, dict
            Default is None
            metadata for the proposal engine
        storage, OPTIONAL, NetCDFStorageView
            Default is None
            If specified, write statistics to this storage layer
        always_change, OPTIONAL, bool
            Default is True
            Currently not implemented -- will always behave as True
            The proposal will always be from the current "state" to the other
            Self proposals will never be made
    """
    def __init__(self, system_generator, residue_name="MOL", atom_expr=None, bond_expr=None, proposal_metadata=None, storage=None, always_change=True):
        super(ButaneProposalEngine, self).__init__(system_generator, residue_name=residue_name, atom_expr=atom_expr, bond_expr=bond_expr, proposal_metadata=proposal_metadata, storage=storage, always_change=always_change)
        self._fake_states = ["butane-A", "butane-B"]
        self.smiles = 'CCCC'

    def _make_skewed_atom_map_old(self, topology):
        """
        Custom definition for the atom map between butane and butane

        (disabled)

        MAP:

                   H - C - CH- C - C - H
          H - C - CH - C - CH- H

        If a regular atom map was constructed (via oechem.OEMCSSearch), all
        atoms would be matched, and the geometry engine would have nothing
        to add.  This method manually finds two of the four carbons and
        matches them to each other, rotated 180degrees.  The other two carbons
        and hydrogens will have to be repositioned by the geometry engine.

        The two carbons are chosen by finding the first carbon-carbon bond in
        the topology. Because geometry needs at least 3 atoms in the atom_map,
        one H from each carbon is also found and mapped to the other.

        Arguments:
        ----------
        topology : app.Topology object
            topology of butane
            Only one topology is needed, because current and proposed are
            identical
        Returns:
        --------
        atom_map : dict
            maps the atom indices of 2 carbons to each other
        """
        for bond in topology.bonds():
            if all([atom.element == app.element.carbon for atom in bond]):
                ccbond = bond
                break
        for bond in topology.bonds():
            if ccbond[0] in bond and any([atom.element == app.element.hydrogen for atom in bond]):
                hydrogen0 = [atom for atom in bond if atom.element == app.element.hydrogen][0]
            if ccbond[1] in bond and any([atom.element == app.element.hydrogen for atom in bond]):
                hydrogen1 = [atom for atom in bond if atom.element == app.element.hydrogen][0]

        atom_map = {
            ccbond[0].index : ccbond[1].index,
            ccbond[1].index : ccbond[0].index,
            hydrogen0.index : hydrogen1.index,
            hydrogen1.index : hydrogen0.index,
        }
        return atom_map

    def _make_skewed_atom_map(self, topology):
        """
        Custom definition for the atom map between butane and butane

        MAP:

        C - C - C - C
            C - C - C - C

        Arguments:
        ----------
        topology : app.Topology object
            topology of butane
            Only one topology is needed, because current and proposed are
            identical
        Returns:
        --------
        atom_map : dict
            maps the atom indices of 2 carbons to each other
        """

        carbons = [ atom for atom in topology.atoms() if (atom.element == app.element.carbon) ]
        neighbors = { carbon : set() for carbon in carbons }
        for (atom1,atom2) in topology.bonds():
            if (atom1.element == app.element.carbon) and (atom2.element == app.element.carbon):
                neighbors[atom1].add(atom2)
                neighbors[atom2].add(atom1)
        end_carbons = list()
        for carbon in carbons:
            if len(neighbors[carbon]) == 1:
                end_carbons.append(carbon)

        # Extract linear chain of carbons
        carbon_chain = list()
        last_carbon = end_carbons[0]
        carbon_chain.append(last_carbon)
        finished = False
        while (not finished):
            next_carbons = list(neighbors[last_carbon].difference(carbon_chain))
            if len(next_carbons) == 0:
                finished = True
            else:
                carbon_chain.append(next_carbons[0])
                last_carbon = next_carbons[0]

        atom_map = {
            carbon_chain[0].index : carbon_chain[2].index,
            carbon_chain[1].index : carbon_chain[1].index,
            carbon_chain[2].index : carbon_chain[0].index,
        }
        return atom_map

class PropaneProposalEngine(NullProposalEngine):
    """
    Custom ProposalEngine to use with PropaneTestSystem defines two "states"
    of butane, identified 'propane-A' and 'propane-B', which are
    tracked by adding a custom _state_key attribute to the topology

    Generates TopologyProposal from propane to propane, only matching
    one CH3 and the middle C such that geometry must rebuild the other

    Can only be used with input topology of propane

    Constructor Arguments:
        system_generator, SystemGenerator object
            SystemGenerator initialized with the appropriate forcefields
        residue_name, OPTIONAL,  str
            Default = "MOL"
            The name that will be used for small molecule residues in the topology
        atom_expr, OPTIONAL, oechem.OEExprOpts
            Default is None
            Currently not implemented -- would dictate how match is defined
        bond_expr, OPTIONAL, oechem.OEExprOpts
            Default is None
            Currently not implemented -- would dictate how match is defined
        proposal_metadata, OPTIONAL, dict
            Default is None
            metadata for the proposal engine
        storage, OPTIONAL, NetCDFStorageView
            Default is None
            If specified, write statistics to this storage layer
        always_change, OPTIONAL, bool
            Default is True
            Currently not implemented -- will always behave as True
            The proposal will always be from the current "state" to the other
            Self proposals will never be made
    """
    def __init__(self, system_generator, residue_name="MOL", atom_expr=None, bond_expr=None, proposal_metadata=None, storage=None, always_change=True):
        super(PropaneProposalEngine, self).__init__(system_generator, residue_name=residue_name, atom_expr=atom_expr, bond_expr=bond_expr, proposal_metadata=proposal_metadata, storage=storage, always_change=always_change)
        self._fake_states = ["propane-A", "propane-B"]
        self.smiles = 'CCC'

    def _make_skewed_atom_map(self, topology):
        """
        Custom definition for the atom map between propane and propane

        If a regular atom map was constructed (via oechem.OEMCSSearch), all
        atoms would be matched, and the geometry engine would have nothing
        to add.  This method manually finds CH3-CH2 and matches each atom to
        itself.  The other carbon and three hydrogens will have to be
        repositioned by the geometry engine.

        The two carbons are chosen by finding the first carbon-carbon bond in
        the topology. To minimize the atoms being rebuilt by geometry, all
        hydrogens from each of the selected carbons are also found and
        mapped to themselves.

        Arguments:
        ----------
        topology : app.Topology object
            topology of propane
            Only one topology is needed, because current and proposed are
            identical
        Returns:
        --------
        atom_map : dict
            maps the atom indices of CH3-CH2 to themselves
        """
        atom_map = dict()
        for bond in topology.bonds():
            if all([atom.element == app.element.carbon for atom in bond]):
                ccbond = bond
                break
        for bond in topology.bonds():
            if any([carbon in bond for carbon in ccbond]) and app.element.hydrogen in [atom.element for atom in bond]:
                atom_map[bond[0].index] = bond[0].index
                atom_map[bond[1].index] = bond[1].index
        return atom_map
