"""
A Python interface to Prime residue interaction data
Copyright Schrodinger, LLC. All rights reserved.
"""
from schrodinger import structure
from schrodinger.infra import mm
[docs]class InteractionReader(object):
    """
    Read the Prime residue interaction data from a structure
    :ivar _available_res: A set of residue names for which Prime interaction
        data has been read
    :vartype _available_res: set
    :ivar _energy: A dictionary of the interaction energy, stored as
        [energy type][res1][res2].  Note that this dictionary is not symmetric:
        [energy type][res1][res2] is stored but [energy type][res2][res1] is not.
        In `_buildEnergyDict` and `getEnergy`, res1 is arbitrarily assigned to be
        less than res2.
    :vartype _energy: dict
    """
    BLOCK_NAME = "m_psp_residue_interaction_energies"
[docs]    def __init__(self, struc):
        """
        Read the Prime residue interaction data
        :param struc: The structure to read interaction data from
        :type struc: `schrodinger.structure.Structure`
        :raise ValueError: If the structure does not have any Prime residue
            interaction data
        """
        # According to Ken Borelli, this copy is necessary due to an issue with
        # the m2io infrastructure
        struc = struc.copy()
        data_names, ninteract, data_handle = self._readEnergyMetaData(struc)
        available_res, energy = self._readEnergyData(data_names, ninteract,
                                                     data_handle)
        self._available_res = available_res
        self._energy = energy 
[docs]    @classmethod
    def hasData(cls, struc):
        """
        Check if the specified structure contains interaction data
        :param struc: The structure to check
        :type struc: `schrodinger.structure.Structure`
        :return: True if the structure has interaction data.  False otherwise.
        :rtype: bool
        """
        struc = struc.copy()
        try:
            cls._openDataHandle(struc)
        except ValueError:
            return False
        else:
            return True 
    def _readEnergyMetaData(self, struc):
        """
        Get metadata about the stored interaction data
        :param struc: The structure to read interaction data from
        :type struc: `schrodinger.structure.Structure`
        :return: A tuple of
              - The names of the stored data (i.e. column headers) (list)
              - The number of interactions (i.e. row count) (int)
              - An open m2io handle to the data (int)
        :rtype: tuple
        """
        data_handle = self._openDataHandle(struc)
        mm.m2io_goto_block(data_handle, self.BLOCK_NAME, 1)
        ninteract = mm.m2io_get_index_dimension(data_handle)
        data_names = mm.m2io_unrequested_data_names(data_handle,
                                                    mm.M2IO_ALL_TYPES)
        return data_names, ninteract, data_handle
    @classmethod
    def _openDataHandle(cls, struc):
        """
        Open an m2io handle to the interaction data
        :param struc: The structure to read interaction data from
        :type struc: `schrodinger.structure.Structure`
        :return: An open m2io handle to the data
        :rtype: int
        :raise ValueError: If the structure does not have any Prime residue
            interaction data
        """
        try:
            data_handle = mm.mmct_ct_m2io_get_unrequested_handle(struc)
        except mm.MmException:
            data_handle = mm.mmct_ct_get_or_open_additional_data(struc, True)
        num_block = mm.m2io_get_number_blocks(data_handle, cls.BLOCK_NAME)
        if num_block == 0:
            err = ("Structure %s does not have Prime residue interaction data" %
                   struc.title)
            raise ValueError(err)
        return data_handle
    def _readEnergyData(self, data_names, ninteract, data_handle):
        """
        Read in the interaction data
        :param data_names: The names of the stored data (i.e. column headers)
        :type data_names: list
        :param ninteract: The number of interactions (i.e. row count)
        :type ninteract: int
        :param data_handle: An open m2iohandle to the data
        :type data_handle: int
        :return: A tuple of (1) The residues for which data was read (set)
            (2) The interaction data (dict).  See `InteractionReader._energy`
            documentation for an explanation of the data format.
        :rtype: tuple
        """
        # The first two data names are for the residues, not energy types
        energy_names = ["r%s" % real_name[1:] for real_name in data_names[2:]]
        available_res = set()
        energy = {name: {} for name in energy_names}
        for i in range(1, ninteract + 1):
            output = mm.m2io_get_string_indexed(data_handle, i, data_names)
            res1 = output.pop(0).strip()
            res2 = output.pop(0).strip()
            available_res.update({res1, res2})
            if res1 > res2:
                res1, res2 = res2, res1
            output = list(map(float, output))
            for cur_energy, cur_name in zip(output, energy_names):
                # Since the values are very sparse, don't bother to store zeroes
                if cur_energy == 0.0:
                    continue
                energy[cur_name].setdefault(res1, {})[res2] = cur_energy
        return available_res, energy
[docs]    def getEnergy(self, energy_type, res1, res2):
        """
        Get the interaction energy between the specified residues
        :param energy_type: The name of the energy to retrieve, such as
            "r_psp_Prime_Energy" or "r_psp_Prime_Covalent"
        :type energy_type: str
        :param res1: The first residue or residue key (i.e. the string value
            returned by `resKey`)
        :type res1: `schrodinger.structure._Residue` or str
        :param res2: The second residue or residue key (i.e. the string value
            returned by `resKey`)
        :type res2: `schrodinger.structure._Residue` or str
        :return: The interaction energy between the specified residues
        :rtype: float
        :raise ValueError: If no data is found for the given energy type or
            residues
        """
        if energy_type not in list(self._energy):
            err = "Energy type %s not found" % energy_type
            raise ValueError(err)
        res1_key = self._getResKey(res1)
        res2_key = self._getResKey(res2)
        res_err = "No interaction energies found for residue %s"
        if res1_key not in self._available_res:
            raise ValueError(res_err % res1_key)
        elif res2_key not in self._available_res:
            raise ValueError(res_err % res2_key)
        if res1_key > res2_key:
            res1_key, res2_key = res2_key, res1_key
        try:
            return self._energy[energy_type][res1_key][res2_key]
        except KeyError:
            return 0.0 
    def _getResKey(self, res):
        """
        Get the Prime-formatted residue name for the specified residue
        :param res: The residue to get the name for or the residue key
        :type res: `schrodinger.structure._Residue` or str
        :return: The formatted name
        :rtype: str
        """
        if isinstance(res, structure._Residue):
            return self.resKey(res)
        else:
            return res.strip()
[docs]    @staticmethod
    def resKey(res):
        """
        Get the Prime-formatted residue name for the specified residue
        :param res: The residue to get the name for
        :type res: `schrodinger.structure._Residue`
        :return: The formatted name
        :rtype: str
        """
        chain = res.chain
        if chain == " ":
            chain = "_"
        resname = res.pdbres.strip()
        inscode = res.inscode.strip()
        return "%s:%s_%i%s" % (chain, resname, res.resnum, inscode)