"""
A module for detecting pi-pi and pi-cation interactions.
Code example for pi-cation interactions::
    from schrodinger.structutils import interactions
    from schrodinger import structure
    recep = None
    for struct in structure.StructureReader(input_file):
        if not recep:
            recep = struct
        else:
            picats = interactions.find_pi_cation_interactions(recep,
                                        struct2=struct)
Code example for pi-pi interactions::
    from schrodinger.structutils import interactions
    from schrodinger import structure
    recep = None
    for struct in structure.StructureReader(input_file):
        if not recep:
            recep = struct
        else:
            pipi = interactions.find_pi_pi_interactions(recep, struct2=struct)
"""
import math
import warnings
import numpy
import schrodinger.infra.structure as infrastructure
import schrodinger.structutils.measure as measure
from schrodinger.infra.mmbitset import Bitset
from schrodinger.structutils import pbc_tools
# Copyright Schrodinger, LLC. All rights reserved.
[docs]def squared_centroid_distance(cent1, cent2):
    """
    Compute the squared distance between two centroids.  Using the squared
    distance saves the sqrt compute cycles.
    :type cent1: Centroid object
    :param cent1: first centroid
    :type cent2: Centroid object
    :param cent2: second centroid
    :rtype: float
    :return: the squared distance between the two centroids.
    """
    vals = [(cent1.xyz[i] - cent2.xyz[i])**2 for i in range(3)]
    return sum(vals) 
[docs]def unit_vector_points(a, b):
    # Finds the unit vector between two points, repeats code from
    # unit_vector just to avoid another function call.
    vec = [(b[x] - a[x]) for x in range(3)]
    mag = math.sqrt(sum([x**2 for x in vec]))
    try:
        return [x / mag for x in vec]
    except ZeroDivisionError:
        return [0.0, 0.0, 0.0] 
[docs]def unit_vector(vec):
    # Returns vec normalized to a unit vector
    mag = math.sqrt(sum([x**2 for x in vec]))
    try:
        return [x / mag for x in vec]
    except ZeroDivisionError:
        return [0.0, 0.0, 0.0] 
[docs]class Centroid(object):
    """
    The object that stores data about a centroid
    """
[docs]    def __init__(self, atoms, xyz):
        """
        Create a Centroid object
        :type atoms: list of int
        :param atoms: the atom numbers involved in the centroid
        :type xyz: list of float
        :param xyz: the XYZ coordinates of the centroid
        """
        self.atoms = atoms
        self.xyz = xyz  
[docs]class CationPiInteraction(object):
    """
    The object that stores the data for a Cation-Pi interaction
    """
[docs]    def __init__(self, cation_structure, pi_structure, cation_centroid,
                 pi_centroid):
        """
        Create a CationPiInteraction object
        :type cation_structure: schrodinger.structure.Structure object
        :param cation_structure: structure that contains the cation
        :type pi_structure: schrodinger.structure.Structure object
        :param pi_structure: structure that contains the pi group
        :type cation_centroid: Centroid object
        :param cation_centroid: Centroid of the positive charge
        :type pi_centroid: Centroid object
        :param pi_centroid: Centroid of the pi group
        :type distance: float
        :param distance: distance between the centroids
        :type angle: float
        :param angle: angle in degrees between the centroids
        """
        self.cation_structure = cation_structure
        """:ivar: structure that contains the cation
           :type: schrodinger.structure.Structure object"""
        self.pi_structure = pi_structure
        """:ivar: structure that contains the pi group
           :type: schrodinger.structure.Structure object"""
        self.cation_centroid = cation_centroid
        """:ivar: Centroid of the positive charge
           :type: Centroid object"""
        self.pi_centroid = pi_centroid
        """:ivar: Centroid of the pi group
           :type: Centroid object""" 
    @property
    def distance(self):
        """
        :ivar: distance between the centroids
        :type: float
        """
        return math.sqrt(
            squared_centroid_distance(self.pi_centroid, self.cation_centroid))
    @property
    def angle(self):
        """
        :ivar: angle in degrees between the centroids
        :type: float
        """
        atom1 = self.pi_structure.atom[self.pi_centroid.atoms[0]]
        atom2 = self.pi_structure.atom[self.pi_centroid.atoms[1]]
        uvec1 = unit_vector_points(self.pi_centroid.xyz, atom1.xyz)
        uvec2 = unit_vector_points(self.pi_centroid.xyz, atom2.xyz)
        # Calculate the norm vector to the plane of the ring.
        ring_plane_norm = numpy.cross(uvec1, uvec2)
        ring_plane_unit_norm = unit_vector(ring_plane_norm)
        # Get the vector between the ring centroid and positive center.
        posring_uvec = unit_vector_points(self.pi_centroid.xyz,
                                          self.cation_centroid.xyz)
        # Calculate the dot product between the ring plane unit vector
        # and positive atom unit
        # vector. From this, the angle can be generated.
        x = numpy.array(ring_plane_unit_norm)
        y = numpy.array(posring_uvec)
        dot_product = numpy.dot(x, y)
        # Make sure rounding error doesn't cause a domain issue
        dot_product = min(max(-1.0, dot_product), 1.0)
        return math.degrees(math.acos(dot_product)) 
[docs]class PiPiInteraction(object):
    """
    The object that stores the data for a Pi-Pi interaction
    """
[docs]    def __init__(self, struct1, struct2, ring1, ring2, face_to_face):
        """
        Create a PiPiInteraction object
        :type struct1: schrodinger.structure.Structure object
        :param struct1: structure that contains the first ring
        :type struct2: schrodinger.structure.Structure object
        :param struct2: structure that contains the second ring
        :type ring1: Centroid object
        :param ring1: Centroid of the first ring
        :type ring2: Centroid object
        :param ring2: Centroid of the second ring
        :type distance: float
        :param distance: distance between the centroids
        :type angle: float
        :param angle: angle in degrees between the two ring planes
        :type face_to_face: bool
        :param face_to_face: True if the interaction is face to face, False if
            it is edge to face
        """
        self.struct1 = struct1
        self.struct2 = struct2
        self.ring1 = ring1
        self.ring2 = ring2
        self.face_to_face = face_to_face 
    @property
    def distance(self):
        """
        :return: distance between the centroids
        :rtype: float
        """
        return math.sqrt(squared_centroid_distance(self.ring1, self.ring2))
    @property
    def angle(self):
        """
        :return: angle in degrees between the centroids
        :rtype: float
        """
        at1 = self.struct2.atom[self.ring2.atoms[0]]
        at2 = self.struct2.atom[self.ring2.atoms[1]]
        at3 = self.struct2.atom[self.ring2.atoms[-1]]
        at4 = self.struct1.atom[self.ring1.atoms[0]]
        at5 = self.struct1.atom[self.ring1.atoms[1]]
        at6 = self.struct1.atom[self.ring1.atoms[-1]]
        return measure.measure_plane_angle(at1,
                                           at2,
                                           at3,
                                           at4,
                                           at5,
                                           at6,
                                           minangle=True) 
[docs]def find_pi_cation_interactions(struct1,
                                struct2=None,
                                atoms1=None,
                                atoms2=None,
                                skip_unknown=None,
                                params=None,
                                honor_pbc=True):
    """
    Determine if any positive centers are within a specified distance cutoff of
    any aromatic ring centroids. For those positive center/ring contacts,
    determine if the positive center is located on the face of the ring rather
    than the edge.
    Code example::
        import interactions
        from schrodinger import structure
        recep = None
        for struct in structure.StructureReader(input_file):
            if not recep:
                recep = struct
            else:
                picats = interactions.find_pi_cation_interactions(recep,
                                            struct2=struct)
    :type struct1: schrodinger.structure.Structure object
    :param struct1: first structure to compute pi-cation interactions for
    :type struct2: schrodinger.structure.Structure object
    :param struct2: second structure to compute pi-cation interactions for or
        or None if the first structure should be used.
    :type atoms1: list of atom indices
    :param atoms1: atoms in struct1 defining the selection to be examined.
        If not passed, all atoms will be used.
    :type atoms2: list of atom indices
    :param atoms2: atoms in struct2 defining the selection to be  examined.
        If not passed, all atoms will be used. If struct2 is not specified
        struct1 will be used.
    :param skip_unknown: UNUSED and deprecated.
    :type params: schrodinger.infra.structure.PiCationParams or NoneType
    :param params: Detailed parameters controlling the interaction criteria.
            If None, then the default values will be used.
    :rtype: list
    :return: list of CationPiInteraction interaction objects::
            # CationPiInteraction properties:
            cation_structure: structure that contains the cation
            pi_structure: structure that contains the pi group
            cation_centroid: Centroid of the positive charge
            pi_centroid: Centroid of the pi group
            distance: distance between the centroids
            angle: angle in degrees between the centroids
    """
    if skip_unknown is not None:
        msg = "The skip_unknown parameter is no longer used."
        warnings.warn(msg, DeprecationWarning, stacklevel=2)
    if params is None:
        params = infrastructure.PiCationParams()
    elif not isinstance(params, infrastructure.PiCationParams):
        msg = ('params must be a schrodinger.infra.structure.PiPiParams '
               'object (found {!r})'.format(params))
        raise TypeError(msg)
    def get_rings_and_atom_bitset(st, atoms=None):
        if atoms is None:
            atoms_bs = Bitset(size=st.atom_total)
            atoms_bs.fill()
            rings = infrastructure.get_sssr(st)
        else:
            atoms_bs = Bitset.from_list(st.atom_total, atoms)
            rings = infrastructure.get_sssr(st, atoms_bs)
        return atoms_bs, rings
    struct2 = struct2 or struct1
    struct1_bs, rings1 = get_rings_and_atom_bitset(struct1, atoms1)
    struct2_bs, rings2 = get_rings_and_atom_bitset(struct2, atoms2)
    pbc = pbc_tools.get_pbc(struct1, struct2, honor_pbc)
    cppinteractions = infrastructure.get_pi_cation_interactions(
        struct1, struct1_bs, rings1, struct2, struct2_bs, rings2, params, pbc)
    pi_cation_interactions = []
    for interaction in cppinteractions:
        cation_atom = interaction.getCation()
        cation = Centroid(
            [cation_atom.getIndex()],
            [cation_atom.x(), cation_atom.y(),
             cation_atom.z()])
        if int(struct1) == cation_atom.getStructure().getHandle():
            pos_structure = struct1
            ring_structure = struct2
        else:
            ring_structure = struct1
            pos_structure = struct2
        pigroup = _get_centroid_from_cpp_structure_ring(interaction.getRing())
        pi_cation_interactions.append(
            CationPiInteraction(pos_structure, ring_structure, cation, pigroup))
    return pi_cation_interactions 
def _get_centroid_from_cpp_structure_ring(ring):
    ring_atoms = [a.getIndex() for a in ring.getAtoms()]
    ring_centroid = [ring.getCenterX(), ring.getCenterY(), ring.getCenterZ()]
    return Centroid(ring_atoms, ring_centroid)
[docs]def gather_rings(astruct, atom_subset=None):
    if atom_subset:
        return infrastructure.get_sssr(astruct, atom_subset)
    else:
        return infrastructure.get_sssr(astruct) 
[docs]def find_pi_pi_interactions(struct1,
                            struct2=None,
                            atoms1=None,
                            atoms2=None,
                            params=None,
                            honor_pbc=True):
    """
    Find all pi-pi interactions between the rings in struct1 and struct2 (or
    within struct1 if only 1 structure is given).  Interactions are classified
    as to whether they are face-to-face or edge-to-face.
    Code example::
        import interactions
        from schrodinger import structure
        recep = None
        for struct in structure.StructureReader(input_file):
            if not recep:
                recep = struct
            else:
                pipi = interactions.find_pi_pi_interactions(recep, struct2=struct)
    :type struct1: schrodinger.structure.Structure object
    :param struct1: first of two structures to compute pi-pi interactions for
    :type struct2: schrodinger.structure.Structure object
    :param struct2: second of two structures to compute pi-pi interactions for.
                    If not given, struct1 will be searched for intramolecular interactions.
    :type atoms1: list of atom indices
    :param atoms1: atoms in struct1 defining the rings that will be examined.
        If not passed, all atoms will be used.
    :type atoms2: list of atom indices
    :param atoms2: atoms in struct2 defining the rings that will be examined.
        If not passed, all atoms will be used. struct2 must be specified if
        this argument is used.
    :type params: schrodinger.infra.structure.PiPiParams or NoneType
    :param params: Detailed parameters controlling the interaction criteria.
                   If None, then the default values will be used.
    :rtype: list of PiPiInteraction objects
    :return: a PiPiInteraction object for each interaction found.
    """
    if params is None:
        params = infrastructure.PiPiParams()
    elif not isinstance(params, infrastructure.PiPiParams):
        msg = ('params must be a schrodinger.infra.structure.PiPiParams '
               'object (found {!r})'.format(params))
        raise TypeError(msg)
    if atoms1 is None:
        rings1 = infrastructure.get_sssr(struct1)
    else:
        bs1 = Bitset.from_list(struct1.atom_total, atoms1)
        rings1 = infrastructure.get_sssr(struct1, bs1)
    if atoms2 is None:
        rings2 = infrastructure.get_sssr(struct2) if struct2 else rings1
    else:
        struct = struct2 or struct1
        bs2 = Bitset.from_list(struct.atom_total, atoms2)
        rings2 = infrastructure.get_sssr(struct, bs2)
    if struct2 is None:
        struct2 = struct1
    pbc = pbc_tools.get_pbc(struct1, struct2, honor_pbc)
    cppinteractions = infrastructure.get_pi_pi_interactions(
        rings1, rings2, params, pbc)
    pi_interactions = []
    for interaction in cppinteractions:
        ring1 = interaction.getRing1()
        ring2 = interaction.getRing2()
        pi_interactions.append(
            PiPiInteraction(struct1, struct2,
                            _get_centroid_from_cpp_structure_ring(ring1),
                            _get_centroid_from_cpp_structure_ring(ring2),
                            interaction.isFaceToFace()))
    return pi_interactions