#!/user/bin/env python
'''structureToAtomInteraction.py

Finds interactions that match the criteria specified by the InteractionFilter

Authorship information:
    __author__ = "Mars (Shih-Cheng) Huang"
    __maintainer__ = "Mars (Shih-Cheng) Huang"
    __email__ = "marshuang80@gmail.com"
    __version__ = "0.2.0"
    __status__ = "done"
'''

from mmtfPyspark.interactions import InteractionFilter, AtomInteraction, InteractionCenter
from mmtfPyspark.utils import ColumnarStructureX
from mmtfPyspark.utils import DistanceBox
import numpy as np


class StructureToAtomInteractions(object):
    '''Class that finds structure to atom intteractions.

    Attributes:
        bfilter (Class): Specifies the conditions for calculating interactions
        pairwise (bool): If True, results as one row per pair interactions.
                         If False, the interactions of one atom with all other
                         atoms are returned as a single row.
    '''

    def __init__(self, bfilter, pairwise=False):

        self.filter = bfilter.value
        self.pairwise = pairwise

    def __call__(self, t):

        interactions = []
        structureId = t[0]

        # convert structure to an array-based format for efficient processing
        arrays = ColumnarStructureX(t[1], True)

        # create a list of query atoms for which interactions should be calculated
        queryAtomIndices = self._get_query_atom_indices(arrays)

        if len(queryAtomIndices) == 0:
            return interactions

        # Add atom (indices) on grid for rapid indexing of atom neighbors on a
        # grid based on a cutoff distance
        box = self._get_distance_box(arrays)

        for queryAtomIndex in queryAtomIndices:
            # find interactions of query atom specified by atom index
            interaction = self._get_interactions(arrays, queryAtomIndex, box)
            interaction.set_structure_id(structureId)

            # only add interations that are within the given limits of interations
            if interaction.get_num_interactions() >= self.filter.get_min_interactions() \
                    and interaction.get_num_interactions() <= self.filter.get_max_interactions():

                # return interactions as either pairs or all interaction of
                # one atom as a row
                if self.pairwise:
                    interactions += interaction.get_pair_interactions_as_rows()
                else:
                    multiInteract = interaction.get_multiple_interactions_as_row(
                        self.filter.get_max_interactions())
                    interactions += multiInteract

        return interactions

    def _get_interactions(self, arrays, queryAtomIndex, box):
        '''Get the interacting neighbors of an atom in a structure

        Attributes
        ----------
            arrays (columnarStructure): structure in columnarStructure format
            queryAtomIndex (int): the index of the querying atom
            box (distanceBox): the distance box of the query atom

        Returns
        -------
            an AtomInteraction class with interacting neighbors
        '''
        interaction = AtomInteraction()

        # get the x,y,z coordinates of the structure
        x = arrays.get_x_coords()
        y = arrays.get_y_coords()
        z = arrays.get_z_coords()

        # get the query atom coordinates
        qx = x[queryAtomIndex]
        qy = y[queryAtomIndex]
        qz = z[queryAtomIndex]

        # get required information of the columnarStructure
        atomToGroupIndices = arrays.get_atom_to_group_indices()
        occupancies = arrays.get_occupancies()
        normalizedbFactors = arrays.get_normalized_b_factors()
        groupNames = arrays.get_group_names()

        # record query atom info
        queryCenter = InteractionCenter(arrays, queryAtomIndex)
        interaction.set_center(queryCenter)

        # Retrieve atom indices of atoms that lay within grid cubes that are
        # within cutoff distance of the query atom
        cutoffDistanceSq = self.filter.get_distance_cutoff() ** 2

        # Retrieve atom indices of atoms that lay within grid cubes
        # that are within cutoff distance of the query atom
        neighborIndices = box.get_neighbors(np.array([qx, qy, qz]))
        # TEST: flattern neighborIndices
        if type(neighborIndices[0]) == list:
            neighborIndices = [
                n for neighbors in neighborIndices for n in neighbors]

        # determine and record interactions with neighbor atoms
        for neighborIndex in neighborIndices:

            # exclude self interactions with a group
            if atomToGroupIndices[neighborIndex] == atomToGroupIndices[queryAtomIndex]:
                continue

            # check if interaction is within distance cutoff
            dx = qx - x[neighborIndex]
            dy = qy - y[neighborIndex]
            dz = qz - z[neighborIndex]
            distSq = dx * dx + dy * dy + dz * dz

            if distSq <= cutoffDistanceSq:
                # Exclude interactions with undesired groups and
                # atoms with partial occupancy (< 1.0)
                if self.filter.is_prohibited_target_group(groupNames[neighborIndex]) \
                        or self.filter.get_normalized_b_factor_cutoff() < normalizedbFactors[neighborIndex] \
                        or occupancies[neighborIndex] < float(1.0):

                    # return an empty atom interaction
                    return AtomInteraction()

                # add interacting atom info
                neighbor = InteractionCenter(arrays, neighborIndex)
                interaction.add_neighbor(neighbor)

                # terminate early if the number of interactions exceeds limit
                if interaction.get_num_interactions() > self.filter.get_max_interactions():
                    return interaction

        return interaction

    def _get_distance_box(self, arrays):
        '''Add atom indices on grid for rapid indexing of atom neighbors on a
        grid based on a cutoff distance

        Attributes
        ----------
            arrays (columnarStructure): structure in columnarStructure format
        '''

        # Get required data
        x = arrays.get_x_coords()
        y = arrays.get_y_coords()
        z = arrays.get_z_coords()
        elements = arrays.get_elements()
        groupNames = arrays.get_group_names()

        box = DistanceBox(self.filter.get_distance_cutoff())
        for i in range(arrays.get_num_atoms()):
            if self.filter.is_target_group(groupNames[i]) \
                    and self.filter.is_target_element(elements[i]):

                newPoint = np.array([x[i], y[i], z[i]])
                box.add_point(newPoint, i)

        return box

    def _get_query_atom_indices(self, arrays):
        '''Returns a list of indices to query atoms in the structure

        Attributes
        ----------
            arrays (columnarStructure): structure in columnarStructure format
        '''

        # Get required data
        groupNames = arrays.get_group_names()
        elements = arrays.get_elements()
        groupStartIndices = arrays.get_group_to_atom_indices()
        occupancies = arrays.get_occupancies()
        normalizedbFactors = arrays.get_normalized_b_factors()

        # Find atoms that match the query criteria and exlcued atoms with
        # partial occupancy
        indices = []
        for i in range(arrays.get_num_groups()):

            start = groupStartIndices[i]
            end = groupStartIndices[i + 1]

            if self.filter.is_query_group(groupNames[start]):
                indices += [j for j in range(start, end)
                            if self.filter.is_query_element(elements[j])
                            and normalizedbFactors[j] < self.filter.get_normalized_b_factor_cutoff()
                            and occupancies[j] >= 1.0]

        return indices
