Source code for schrodinger.application.matsci.enumeration

"""
A class and function that help enumerate possibilities - could be enumerating
over items of a list or enumerating the element at various atom positions in a
structure, for instance.

Copyright Schrodinger, LLC. All rights reserved.
"""

import argparse
import sys

from schrodinger.application.matsci import buildcomplex
from schrodinger.structutils import analyze
from schrodinger.structutils import build
from schrodinger.utils import cmdline

_version = '$Revision: 0.0 $'


class Alchemist(object):
    """
    Class that mutates the element of an atom in a structure
    """

    def __init__(self, master, targets, new_items, child=None):
        """
        Create an Alchemist object. Alchemists modify an object by changing a
        target item to new item.  For instance, an Alchemist:
        - Might transmute the element of a particular atom
        - Might change the item of a list to a new value
        For enumeration, a series of Alchemists will work together to perform
        all necessary mutations, each alchemist passing their mutated structures
        on to another (child) alchemist. The general workflow looks like::

            receive mutated_object from parent alchemist
            for target in targets:
                for new_item in new_items:
                    make a copy_of_mutated_object
                    change item at target in copy_of_mutated_object to new_item
                    pass copy_of_mutated_object to child alchemist

        :type master: object
        :param master: The master for this object. Can be a panel or other
            object that has methods that get called by Alchemist methods. By
            default, the Alchemist class calls master.transmutationCompleted with
            each fully transmuted object.

        :type targets: list
        :param targets: list of possible positions for this Alchemist to mutate

        :type new_items: list
        :param new_items: list of items to change each target to.

        :type child: `Alchemist`
        :param child: The Alchemist to call each time this Alchemist transmutes
            an atom.  If child is None, than this Alchemist is the last in the line.
        """

        self.master = master
        self.original_targets = targets
        self.new_items = new_items
        self.child = child

    def performTransmutation(self, data, index, item):
        """
        Perform a transmutation on a copy of data at index.  It is very
        important that this happen on a copy of data, not data itself.

        :param data: The object to transmute a copy of

        :param index: An index that tells what part of data to transmute

        :param item: The object to mutate data[index] to

        :rtype: type(data)
        :return: A COPY of data, transmuted at index
        """

        mutated_data = data[:]
        mutated_data[index] = item
        return mutated_data

    def allTransmutationsCompleted(self, mydata, myused):
        """
        Perform whatever action is required when all transmutations have been
        performed on the mydata object

        :param mydata: The object that is fully transmuted

        :type myused: list
        :param myused: List of indexes that have been transmuted in mydata
        """

        self.master.transmutationCompleted(mydata, myused)

    def runTransmutations(self, data, used=None):
        """
        Begin transmutations.

        This Alchemist will determine what targets to transmute based on
        its set of original targets minus those targets that have been used by
        previous Alchemists.  Each time this Alchemist transmutes an item, it
        calls its child Alchemist to do its mutations (which calls its child
        Alchemist each time it transmutes an item, and so on).  Thus is the
        iterative process of transmutating all possible combinations achieved.

        If this Alchemist does not have any children, than it is the last in the
        Alchemist line and it should call allTransmutationsCompleted when it is
        done with a transmutation.

        :type data: object
        :param data: The object to be transmuted

        :type used: list of int
        :param used: The indexes in data of the transmutations performed by
            parent Alchemists.  The last item in this list is used to ensure that
            this Alchemist does not mutate any target already covered by parent
            Alchemists.
        """

        targets = self.original_targets[:]
        if used:
            # Remove any targets that have been already used by ancestors
            try:
                index = targets.index(used[-1])
                targets = targets[index + 1:]
            except ValueError:
                pass
        else:
            # This is the original Alchemist, no targets have been used
            used = []

        # Cycle through all targets, transmuting them one by one
        for index in targets:

            # Cycle through all new items, transmuting to them one by one
            for item in self.new_items:
                # Start with a copy each time so my own trasmutations don't
                # accumulate
                myused = used[:]
                myused.append(index)

                mydata = self.performTransmutation(data, index, item)

                if self.child:
                    # Let my child iterate its transmutations on this structure
                    self.child.runTransmutations(mydata, used=myused)
                else:
                    # This is the last Alchemist in the line, we are done with
                    # this structure
                    self.allTransmutationsCompleted(mydata, myused)


class StructureAlchemist(Alchemist):
    """
    An abstract base Alchemist class for Alchemists that transmute Structures -
    subclasses will probably need to re-implement the performTransmutation and
    allTransmutationsCompleted methods.
    """

    def __init__(self, master, targets, new_items, path, child=None):
        """
        Create an Alchemist object

        :type path: str
        :param path: The path to a temporary file to write structures to

        See parent class for additional documentation
        """

        Alchemist.__init__(self, master, targets, new_items, child=child)
        self.path = path

    def runTransmutations(self, data, used=None):
        """
        Begin transmuating atoms.

        This Alchemist will determine what target atoms to transmute based on
        its set of original targets minus those targets that have been used by
        previous Alchemists.  Each time this Alchemist transmute an atom, it
        calls its child Alchemist to do its mutations (which calls its child
        Alchemist each time it transmutes an atom).  Thus is the iterative
        process of transmutating all possible combinations achieved.

        If this Alchemist does not have any children, than it is the last in the
        Alchemist line and it should add a new project entry each time it
        transmutes an atom.

        :type struct: `schrodinger.structure.Structure`
        :param struct: The structure object with the transmuted atoms

        :type used: list of int
        :param used: The atom indexes of the transmuted atoms by parent
            Alchemists.  The last atom in this list is used to ensure that this
            Alchemist does not mutate any target already covered by parent
            Alchemists.
        """

        Alchemist.runTransmutations(self, data, used=used)
        if not self.child:
            self.master.showProgress()


class ElementAlchemist(StructureAlchemist):
    """
    An alchemist class that transmutes elements in a Structure

    :note: The new_items argument to the constructor for this class should have
        the following format::

            type new_items: list of tuple
            param new_items: Each item of the list is an (element, color) tuple,
            where element is the atomic symbol of the element to mutate to, and
            color is the new color for that atom in a form accepted by the
            _StructAtom.color property
    """

    def performTransmutation(self, struct, index, item):
        """
        Perform a transmutation on a copy of struct at index.  It is very
        important that this happen on a copy of struct, not struct itself.

        :type struct: `schrodinger.structure.Structure`
        :param struct: The structure to transmute a copy of

        :type index: int
        :param index: An index that tells what atom to transmute

        :type item: tuple
        :param item: An (element, color) tuple,
            where element is the atomic symbol of the element to mutate to, and
            color is the new color for that atom in a form accepted by the
            _StructAtom.color property

        :rtype: `schrodinger.structure.Structure`
        :return: A COPY of struct with atom index transmuted
        """

        element, color = item

        # Actually change the element
        mystruct = struct.copy()
        atomobj = mystruct.atom[index]
        buildcomplex.transmute_atom(atomobj, element, color=color)

        # Update the structure's title
        title = f'{struct.title}_{index}:{element}'
        mystruct.title = title

        return mystruct

    def fixHydrogens(self, struct, used):
        """
        Correct the number of hydrogens bonded to each transmuted atom

        :type struct: `schrodinger.structure.Structure`
        :param struct: The structure object with the transmuted atoms

        :type used: list of int
        :param used: The atom indexes of the transmuted atoms
        """

        delatoms = []
        for index in used:
            atomobj = struct.atom[index]
            for neighbor in atomobj.bonded_atoms:
                if neighbor.element == 'H':
                    delatoms.append(neighbor.index)
        atom_map = struct.deleteAtoms(delatoms, renumber_map=True)
        used_atom_map = [atom_map[x] for x in used]
        build.add_hydrogens(struct, atom_list=used_atom_map)

    def allTransmutationsCompleted(self, struct, used):
        """
        All transmutations have been performed on struct, write it out to a
        file if we are keeping this result

        :type struct: `schrodinger.structure.Structure`
        :param struct: The structure object with the transmuted atoms

        :type used: list of int
        :param used: The atom indexes of the transmuted atoms
        """

        self.fixHydrogens(struct, used)
        if self.master.transmutedStructureCompleted(struct):
            struct.retype()
            struct.append(self.path)


class VacancyAlchemist(StructureAlchemist):
    """
    An alchemist class that deletes atoms in a Structure and the hydrogens they
    are bound to

    :note: The new_items argument to the constructor for this class should have
        the following format::

            type new_items: list
            param new_items: Should be [None]
    """

    DELETE_PROPERTY = 'b_matsci_impending_vacancy'

    def performTransmutation(self, struct, index, item):
        """
        Mark atom index for deletion, and also mark all the hydrogen atoms it is
        attached to. This will happen on a copy of the input structure.

        Normally this method in parent classes actually changes the structure,
        but in this case, deleting atoms would mess up all the atom numbering
        for Alchemists down the line, so we only mark those atoms for later
        deletion.

        :type struct: `schrodinger.structure.Structure`
        :param struct: The structure to transmute a copy of

        :type index: int
        :param index: An index that tells what atom to delete

        :type item: None
        :param item: None - unused

        :rtype: `schrodinger.structure.Structure`
        :return: A COPY of struct with atom index and attached hydrogens marked
            for deletion.
        """

        mystruct = struct.copy()
        atom = mystruct.atom[index]
        atom.property[self.DELETE_PROPERTY] = True
        for neighbor in atom.bonded_atoms:
            if neighbor.element == 'H':
                neighbor.property[self.DELETE_PROPERTY] = True
        return mystruct

    def deleteAtoms(self, struct):
        """
        Delete all atoms marked for deletion

        :type struct: `schrodinger.structure.Structure`
        :param struct: The structure object with the transmuted atoms
        """

        delatoms = analyze.evaluate_asl(struct,
                                        'atom.%s' % self.DELETE_PROPERTY)
        struct.deleteAtoms(delatoms)

    def allTransmutationsCompleted(self, struct, used):
        """
        All transmutations have been performed on struct, write it out to a
        file if we are keeping this result

        :type struct: `schrodinger.structure.Structure`
        :param struct: The structure object with the transmuted atoms

        :type used: list of int
        :param used: The atom indexes of the transmuted atoms
        """

        title = '_'.join([struct.title, 'x'] + [str(x) for x in used])
        self.deleteAtoms(struct)
        if self.master.transmutedStructureCompleted(struct):
            struct.retype()
            struct.title = title
            struct.append(self.path)


def perform_transmutations(minimum_transmutations, maximum_transmutations,
                           transmutation_sites, alchemist_creator, data,
                           **kwargs):
    """
    Enumerate all possible transmutations on an object

    :type minimum_transmutations: int
    :param minimum_transmutations: The minimum number of transmutations to
        perform on each object

    :type maximum_transmutations: int
    :param maximum_transmutations: The maximum number of transmutations to
        perform on each object

    :type transmutation_sites: list
    :param transmutation_sites: The indexes of the object that are available for
        transmutation

    :type alchemist_creator: callable
    :param alchemist_creator: A method that creates an Alchemist object. The
        method should take the following positional arguments: 1) A list of sites to
        transmute, 2) A child alchmist (or None if no child). In addition, any
        additional keyword arguments passed to the perform_transmutations function
        are passed on to the alchemist_creator function as keyword arguments.

    :param data: The object to perform transmutations on

    :note: Any additional keyword arguments are passed to the alchemist_creator
        function
    """

    num_sites = len(transmutation_sites)
    for num_transmutations in range(minimum_transmutations,
                                    maximum_transmutations + 1):
        alchemists = []
        # Create the alchemists that will perform the transmutations - note
        # that they are created in order of: newest child -> oldest ancestor
        for mutnum in range(num_transmutations):
            start = num_transmutations - mutnum - 1
            end = num_sites - mutnum
            # Each alchemist is given a slice of the target sites to work on.
            # The oldest (first in line, last created) alchemist will get the
            # first target site plus all sites up to n-m, where n is the number
            # of sites and m is the number of transmutations to perform. The
            # youngest (last in line, first created) will get sites m through n.
            targets = transmutation_sites[start:end]
            try:
                child = alchemists[-1]
            except IndexError:
                child = None
            alchemists.append(alchemist_creator(targets, child, **kwargs))
        # Perform the transmutations by starting up the oldest ancestor
        alchemists[-1].runTransmutations(data)


description = __doc__


def main(*args):
    parser = argparse.ArgumentParser(description=description, add_help=False)
    parser.add_argument('-h',
                        '-help',
                        action='help',
                        default=argparse.SUPPRESS,
                        help='Show this help message and exit.')
    parser.add_argument('-v',
                        '-version',
                        action='version',
                        default=argparse.SUPPRESS,
                        version=_version,
                        help="Show the program's version number and exit.")
    parser.parse_args(args)


# For running outside of Maestro:
if __name__ == '__main__':
    cmdline.main_wrapper(main, *sys.argv[1:])