import collections
from schrodinger import structure
[docs]class LDData:
    """
    Class for storing information about Maestro-level data that may be exported
    to LiveDesign. (Objects from this class should only store Maestro-level
    data, not LiveDesign data.)
    Specifically, this class is not meant to store the values of the data to be
    exported, but rather to serve as a "view" to that data. It should:
    1. provide identifying information to the user about the data that is
       available for export (e.g. by providing user-friendly property names to
       tables and other view widgets), and
    2. provide an internal "pointer" to data to be exported (by uniquely
       identifying exactly what data should be exported, as selected by the
       user)
    This class is meant to be general enough to store information about both
    structure property data and attachment data (e.g. 3D, images, files).
    """
[docs]    def __init__(self,
                 data_name=None,
                 family_name=None,
                 user_name=None,
                 requires_3d=False,
                 requires_ffc=False):
        """
        Instances representing structure properties should be initialized with
        the property data name (e.g. `data_name="s_m_title"`).
        If there is no data name, e.g. for image and attachment data, instances
        should be initialized with both the family name and the user name (e.g.
        "Atom Properties", "Hot Atoms").
        If the family or user name is supplied along with the data name, these
        will be used for display purposes within the panel rather than the
        values derived from the data name.
        :raise ValueError: if insufficient information is supplied on
            initialization to fully specify a unique instance
        """
        self._requires_3d = requires_3d
        self._requires_ffc = requires_ffc
        # If not provided on initialization, these values will be determined
        # from the data name
        self._family_name = family_name
        self._user_name = user_name
        if not data_name and not (family_name and user_name):
            msg = ('LDData requires either a structure property data name or'
                   ' both family name and user name.')
            raise ValueError(msg)
        elif data_name:
            self._prop_name = structure.PropertyName(data_name)
        else:
            self._prop_name = None 
    @property
    def family_name(self):
        """
        Return the family name if provided on initialization. Otherwise,
        determine family name from the data name.
        :return: family name for display
        :rtype: str
        """
        if self._family_name:
            return self._family_name
        if self._prop_name:
            return get_long_family_name(self._prop_name.family)
    @property
    def user_name(self):
        """
        Return the user name if provided on initialization. Otherwise, determine
        user name from the data name.
        :return: user name for display
        :rtype: str
        """
        if self._user_name:
            return self._user_name
        if self._prop_name:
            return self._prop_name.userName()
    @property
    def data_name(self):
        """
        Return data name if provided on initialization.
        :return: data name or `None`
        :rtype: `str` or `None`
        """
        if self._prop_name:
            return self._prop_name.dataName()
        return None
    @property
    def requires_3d(self):
        """
        :return: whether the export of this data requires the export of 3D data
        :rtype: bool
        """
        return self._requires_3d
    @property
    def requires_ffc(self):
        """
        :return: whether this data needs to be export as a freeform column
        :rtype: bool
        """
        return self._requires_ffc
    def __repr__(self):
        return str(self)
    def __str__(self):
        msg = 'LDData(data_name="{0}", family_name="{1}", user_name="{2}")'
        return msg.format(self.data_name, self.family_name, self.user_name)
    def __eq__(self, other):
        if not isinstance(other, LDData):
            return False
        return (self.family_name == other.family_name and
                self.user_name == other.user_name and
                self.data_name == other.data_name)
    def __hash__(self):
        return hash((self.data_name, self.family_name, self.user_name)) 
[docs]class ReceptorLigandPair:
    """
    Data class for storing a receptor structure and a ligand structure.
    """
[docs]    def __init__(self, receptor=None, ligand=None):
        self.receptor = receptor
        self.ligand = ligand 
    def __copy__(self):
        rec = self.receptor.copy() if self.receptor else None
        lig = self.ligand.copy() if self.ligand else None
        return ReceptorLigandPair(rec, lig)
    def __hash__(self):
        return hash((self.receptor, self.ligand))
    def __eq__(self, other):
        if not isinstance(other, type(self)):
            return False
        return (self.receptor, self.ligand) == (other.receptor, other.ligand) 
[docs]class ReceptorLigandGroup:
    """
    Data class for unambiguously storing a group of receptor and ligand
    structures. In addition to the primary ligand, this class also supports
    storing an "alternate" ligand meant for 3D upload in place of the primary
    ligand and "additional" ligands meant for 3D upload in addition to the
    primary ligand.
    The alternate ligand may have a distinct ligand pose, or it may have
    a different structure entirely. For example, the primary ligand may be a
    ligand after covalently binding to a receptor, and the alternate ligand may
    be the independent ligand molecule prior to complexing.
    Each of these structures is optional; a `ReceptorLigandGroup` instance
    will not always need to contain a receptor and a ligand, or any other
    structures.
    """
[docs]    def __init__(self,
                 receptor=None,
                 ligand=None,
                 alt_ligand=None,
                 add_rl_pairs=None):
        """
        :param receptor: a receptor structure or None
        :type receptor: structure.Structure or None
        :param ligand: a ligand structure or None
        :type ligand: structure.Structure or None
        :param alt_ligand: extra structure data associated with the ligand, e.g.
            to be used as 3D data for certain LiveDesign uploads
        :type alt_ligand: structure.Structure or None
        :param add_rl_pairs: additional ligand/receptor structures to be
            uploaded as 3D data in addition to the conventional (ligand or
            alternate ligand) 3D uploads
        :type add_rl_pairs: list(ReceptorLigandPair) or None
        """
        self.receptor = receptor
        self.ligand = ligand
        self.alt_ligand = alt_ligand
        self.add_rl_pairs = add_rl_pairs or [] 
    def __copy__(self):
        rec_copy = self.receptor.copy() if self.receptor else None
        lig_copy = self.ligand.copy() if self.ligand else None
        alt_lig_copy = self.alt_ligand.copy() if self.alt_ligand else None
        add_rl_pair_copies = [
            add_rl_pair.copy() for add_rl_pair in self.add_rl_pairs
        ]
        return ReceptorLigandGroup(receptor=rec_copy,
                                   ligand=lig_copy,
                                   alt_ligand=alt_lig_copy,
                                   add_rl_pairs=add_rl_pair_copies) 
[docs]class ReceptorLigandMap(collections.defaultdict):
    """
    A specialized dictionary for organizing receptor and ligand structures. Each
    key points to a list of receptor-ligand groups associated with that key. For
    convenience, this class also features several generators.
    """
[docs]    def __init__(self):
        """
        Initialize as a `list`-based `collections.defaultDict`.
        """
        super(ReceptorLigandMap, self).__init__(list) 
    @property
    def num_rl_groups(self):
        """
        :return: the number of receptor-ligand groups stored in this object.
        :rtype: int
        """
        count = 0
        for item in self.rl_group_items:
            count += 1
        return count
    @property
    def rl_group_items(self):
        """
        :return: an iterator for (key, receptor-ligand group) pairs stored in
            this map. Note that each key can correspond to multiple receptor-
            ligand groups.
        :rtype: iterator((str, ReceptorLigandGroup))
        """
        for key, group_list in self.items():
            for rl_group in group_list:
                yield key, rl_group
    @property
    def rl_groups(self):
        """
        :return: an iterator for receptor-ligand group objects stored in this
            map.
        :rtype: iterator(ReceptorLigandGroup)
        """
        for key, rl_group in self.rl_group_items:
            yield rl_group
    @property
    def ligand_items(self):
        """
        :return: an iterator for (key, ligand) pairs stored in this map
        :rtype: iterator((str, structure.Structure))
        """
        for key, rl_group in self.rl_group_items:
            if rl_group.ligand:
                yield key, rl_group.ligand
    @property
    def ligands(self):
        """
        :return: an iterator for ligand structures stored in this map
        :rtype: iterator(structure.Structure)
        """
        for key, ligand_st in self.ligand_items:
            yield ligand_st
    @property
    def alt_ligand_items(self):
        """
        :return: an iterator for (key, alternative ligand) pairs stored in this
            map
        :rtype: iterator((str, structure.Structure))
        """
        for key, rl_group in self.rl_group_items:
            if rl_group.alt_ligand:
                yield key, rl_group.alt_ligand
    @property
    def alt_ligands(self):
        """
        :return: an iterator for alternative ligand structures stored in this
            map
        :rtype: iterator(structure.Structure)
        """
        for key, alt_lig_st in self.alt_ligand_items:
            yield alt_lig_st
    @property
    def add_rl_pair_items(self):
        """
        :return: an iterator for (key, additional receptor-ligand pair) 2-tuples
            stored in this map
        :rtype: iterator((str, ReceptorLigandPair))
        """
        for key, rl_group in self.rl_group_items:
            for add_rl_pair in rl_group.add_rl_pairs:
                yield key, add_rl_pair
    @property
    def add_rl_pairs(self):
        """
        :return: an iterator for additional receptor-ligand pairs stored on this
            map
        :rtype: iterator(ReceptorLigandPair)
        """
        for key, rl_pair in self.add_rl_pair_items:
            yield rl_pair
    @property
    def receptor_items(self):
        """
        :return: an iterator for (key, receptor) pairs stored in this map
        :rtype: iterator((str, structure.Structure))
        """
        for key, rl_group in self.rl_group_items:
            if rl_group.receptor:
                yield key, rl_group.receptor
    @property
    def receptors(self):
        """
        :return: an iterator for receptor structures stored in this map
        :rtype: iterator(structure.Structure)
        """
        for key, receptor_st in self.receptor_items:
            yield receptor_st
    @property
    def structures(self):
        """
        :return: an iterator for all structures stored in this map
        :rtype: iterator(structure.Structure)
        """
        for rl_group in self.rl_groups:
            if rl_group.receptor is not None:
                yield rl_group.receptor
            if rl_group.ligand is not None:
                yield rl_group.ligand
            for add_rl_pair in rl_group.add_rl_pairs:
                if add_rl_pair.receptor is not None:
                    yield add_rl_pair.receptor
                if add_rl_pair.ligand is not None:
                    yield add_rl_pair.ligand
    def __copy__(self):
        rl_map_copy = ReceptorLigandMap()
        st_copy_cache = {None: None}
        def get_cached_copy(st):
            """
            Retrieve a structure from the cache, or create one and cache it.
            :param st: a structure to copy, or `None`
            :type st: structure.Structure or NoneType
            :return: a copy of `st`, or `None` (if `st` is `None`)
            :rtype: structure.Structure or NoneType
            """
            if st not in st_copy_cache:
                st_copy_cache[st] = st.copy()
            return st_copy_cache[st]
        for key, rl_group in self.rl_group_items:
            rl_group_copy = ReceptorLigandGroup()
            rl_group_copy.receptor = get_cached_copy(rl_group.receptor)
            rl_group_copy.ligand = get_cached_copy(rl_group.ligand)
            rl_group_copy.alt_ligand = get_cached_copy(rl_group.alt_ligand)
            for rl_pair in rl_group.add_rl_pairs:
                rec_copy = get_cached_copy(rl_pair.receptor)
                lig_copy = get_cached_copy(rl_pair.ligand)
                rl_pair_copy = ReceptorLigandPair(receptor=rec_copy,
                                                  ligand=lig_copy)
                rl_group_copy.add_rl_pairs += [rl_pair_copy]
            rl_map_copy[key].append(rl_group_copy)
        return rl_map_copy 
[docs]def get_long_family_name(short_family_name):
    """
    Given the short family name of a structure property (e.g. "m"), return the
    corresponding long family name (e.g. "Maestro"). If no long family name is
    defined, return the short family name argument.
    :param short_family_name: the short family name of a structure property
    :param short_family_name: str
    :return: the corresponding long family name if one exists, otherwise the
        short family name
    :rtype: str
    """
    return structure.PROP_LONG_NAME.get(short_family_name, short_family_name)