"""
sstrajectory.py

sstrajectory is the main class through which simulation trajectories are read in.

"""

##
##                                       _ _              _ 
##   ___ __ _ _ __ ___  _ __   __ _ _ __(_) |_ _ __ __ _ (_)
##  / __/ _` | '_ ` _ \| '_ \ / _` | '__| | __| '__/ _` || |
## | (_| (_| | | | | | | |_) | (_| | |  | | |_| | | (_| || |
##  \___\__,_|_| |_| |_| .__/ \__,_|_|  |_|\__|_|  \__,_|/ |
##                     |_|                             |__/ 
##
##
## Alex Holehouse (Pappu Lab)
## Simulation analysis package
## Copyright 2014 - 2021
##

import mdtraj as md
import numpy as np
from .configs import *
from .ssdata  import ALL_VALID_RESIDUE_NAMES

from .ssprotein import SSProtein
from .ssexceptions import SSException
from . import ssutils
from . import ssio

from copy import copy


class SSTrajectory:
    """
    SSTrajectory class that holds a single simulation trajectory object. 
    

    """

    #oxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxo
    #
    #
    def __init__(self, trajectory_filename=None, pdb_filename=None, TRJ=None, protein_grouping=None, pdblead=False, debug=False, extra_valid_residue_names=None):
        """
        CAMPARITraj trajectory object initializer. 

        CAMPARITraj will, by default, extract out the protein component from 
        your trajectory automatically, which lets you ask questions about the
        protein only (i.e. without salt ions getting in the way).
        
        Note that by default the mechanism by which individual proteins are
        identified is by cycling through the unique chains and determining 
        if they are protein or not. You can also provide manual grouping
        via the protein_grouping option, which lets you define which
        residues should make up an individual protein. This can be useful
        if you have multiple proteins associated with the same chain, which
        happens in CAMPARI if you have more than 26 separate chains (i.e.
        every protein after the 26th is the 'Z' chain).
               
        Parameters
        ----------        
        trajectory_filename : str
            Filename which contains the trajectory file of interest. Normally \
            this is `__traj.xtc` or `__traj.dcd`.

        pdb_filename : str
            Filename which contains the pdb file associated with the trajectory \
            of interest. Normally this is `__START.pdb`.

        TRJ : mdtraj.Trajectory
            It is sometimes useful to re-defined a trajectory and create a new SSTrajectory \
            object from that trajectory. This could be done by writing that new trajectory \
            to file, but this is extremely slow due to the I/O impact of reading/writing \
            from disk. If an mdtraj trajectory objected is passed, this is used as the \
            new trajectory from which the SSTrajectory object is constructed. 

            Default = None

        protein_grouping : list of lists of ints
            Lets you manually define protein groups to be considered independently.
        
            Default = None

        pdblead : bool
            Lets you set the PDB file (which is normally ONLY used as a topology \
            file) to be the first frame of the trajectory. This is useful when \
            the first PDB file holds some specific reference information which \
            you want to use (e.g. RMSD or Q). 

            Default = False

        debug : book
            Prints warning/help information to help debug weird stuff during initial 
            trajectory read-in. Default = False.
            

        extra_valid_residue_names : list
            By default SOURSOP identifies chains as proteins based on the a set of
            normally seen protein residue names. These are defined in soursop/ssdata,
            and are listed below:

            ['ALA','CYS','ASP','ASH','GLU','GLU','PHE','GLY','HIE','HIS','HID','HIP',
             'ILE','LEU', 'LYS','LYD','MET','ASN','PRO','GLN','ARG','SER','THR','VAL',
             'TRP','TYR','AIB', 'ABA','NVA','NLE', 'ORN', 'DAB','PTR','TPO','SEP', 'KAC', 
             'KM1', 'KM2' 'KM3', 'ACE','NME', 'FOR', 'NH2']
        
            This keyword allows the user to pass a list of ADDITIONAL residues that
            we want SOURSOP to recognize as valid residues to extract a chain as a
            protein molecule.
        """

        self.valid_residue_names = []
        self.valid_residue_names.extend(ALL_VALID_RESIDUE_NAMES)

        if extra_valid_residue_names is not None:
            try:
                self.valid_residue_names.extend(extra_valid_residue_names)
            except Exception:
                print('Unable to use the extra_valid_residue_names - this must be a list of strings')
        
        # first we decide if we're reading from file or from an existing trajectory
        if (trajectory_filename is None) and (pdb_filename is None):
            if TRJ is None:
                raise SSException('No input provided! Please provide ether a pdb and trajectory file OR a pre-formed traj object')
                
            # note the [:] means this is a COPY!
            self.traj = TRJ[:]
        else:
            if (trajectory_filename is None):
                raise SSException('No trajectory file provided!')
            if (pdb_filename is None):
                raise SSException('No PDB file provided!')

            # read in the raw trajectory
            self.traj = self.__readTrajectory(trajectory_filename, pdb_filename, pdblead)


        # Next, having read in the trajectory we parse out into proteins
        # extract a list of protein trajectories where each protein is assumed
        # to be in its own chain
        if protein_grouping == None:
            self.proteinTrajectoryList = self.__get_proteins(self.traj, debug)        
        else:
            self.proteinTrajectoryList = self.__get_proteins_by_residue(self.traj, protein_grouping, debug)

        self.__single_protein_traj = self.__get_all_proteins(self.traj)

        
        self.num_proteins = len(self.proteinTrajectoryList)
        self.n_frames = len(self.traj)


    def  __repr__(self):
        return "SSTrajectory (%s): %i proteins and %i frames" % (hex(id(self)), self.num_proteins, self.n_frames)


    def __len__(self):
        # Edited to mimic the behavior of `mdtraj` trajectory objects.
        # Originally: `return (self.num_proteins, self.n_frames)`
        return self.n_frames


    def length(self):
        # Implemented a method that encapsulates the data output by the original `__len__` method.
        return (self.num_proteins, self.n_frames)


    #oxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxo
    #
    #
    def __readTrajectory(self, trajectory_filename, pdb_filename, pdblead):
        """
        Internal function which parses and reads in a CAMPARI trajectory

        Read a trajectory file. This was separated out into its own
        function in case we want to add additional sanity checks during
        the file loading.

        Notably older versions of CAMPARI mess up the unitcell length
        vectors, so will cause problems, but you can get around this by
        having GROMACS rebuild the trajectory. If this happens an error
        pops up but instructions on how to use GROMACS to fix it are
        presented.

        Parameters
        -----------
        trajectory_filename : str
            Filename which contains the trajectory file of interest. File type
            is automatically detected and dealt with mdtraj' 'load' command
            (i.e. md.load(filename, top=pdb_filename))

        pdb_filename : str
            Filename which contains the pdb file associated with the trajectory
            of interest. This defines the topology of the system and must match
            the trajectory in terms of number of atomas


        pdblead : bool
            Also extract the coordinates from the PDB file and append it to 
            the front of the trajectory. This is useful if you are starting
            an analysis where that first structure should be a reference frame
            but it's not actually included in the trajectory file.

        Returns
        --------
        mdtraj.traj 
            Returns an mdtraj trajectory object

        """

        # straight up read the trajectory first using mdtraj's awesome
        # trajectory reading facility
        traj =  md.load(trajectory_filename, top=pdb_filename)
                    
        # check unit cell lengths
        try:
            uc_lengths = traj.unitcell_lengths[0]

            # this is s custom warning for a specific edge-case we encounter a lot
            if (uc_lengths[0] == 0 or uc_lengths[1] == 0 or uc_lengths[2] == 0):
                ssio.warning_message("Trajectory file unit cell lengths are zero for at least one dimension. This is a probably a bug with an FRC generated __START.pdb file, because in the old version of CAMPARI used to do grid based FRC calculations the unit cell dimensions are not written correctly. This may cause issues but we're going to assume everything is OK for now. Check the validity of any analysis output. If you're worried, you can use the following workaround.\n\n:::: WORK AROUNDS ::::\nSimply run\n\ntrjconv -f __traj.xtc -s __START.pdb -box a b c -o frc.xtc \n\nAn then \n\ntrjconv -f frc.xtc -s __START.pdb -box a b c -o start.pdb -dump 0\n\n\nHere\n-f n__traj.xtc   : defines the trajectory file\n-s __start.pdb   : defines the pdb file used to parse the topology\n-box a b c       : defines the box lengths **in nanometers**\n-o frc.xtc       : is the name of the new trajectory file with updated box lengths\nSelect 0 (system) when asked to 'Select group for output'.The second step creates the equivalent PDB file with the header-line correctly defining the box unit cell lengths and angles. These two new files should then be used for analysis.\n\nAs an example, if my FRC simulation had a sphere radius of 100 angstroms then my correction command would look something like \n\ntrjconv -f __traj.xtc -s __START.pdb -box 20 20 20 -o frc.xtc\ntrjconv -f frc.xtc -s __START.pdb -box 20 20 20 -o start.pdb -dump 0")

        except TypeError:
            ssio.warning_message("Warning: UnitCell lengths were not provided... This may cause issues but we're going to assume everything is OK for now...")
        
        # if pdbLead is true then load the pdb_filename as a trajectory
        # and then add it to the front (the PDB file is its own topology
        # file so no need to specificy the top= file here!
        if pdblead:
            pdbtraj = md.load(pdb_filename)
            traj = pdbtraj+traj


            # having added the PDB file now check all the unit-cells match up!
            try:
                uc_lengths_1 = traj.unitcell_lengths[0]
                uc_lengths_2 = traj.unitcell_lengths[1]
            
                if (uc_lengths_1[0] != uc_lengths_2[0]) or (uc_lengths_1[2] != uc_lengths_2[2]) or (uc_lengths_1[2] != uc_lengths_2[2]):
                    ssio.warning_message('........................\nWARNING:\nThe unit cell dimensions of the PDB file and trajectory file did not match, specifically\nPDB file = [ %s ]\nXTC file = [ %s ]\nThis may cause issues if native MDTraj untilities are used (and potentially for SOURSOP utilities that are based on these. It is not necessarily an issue, but PLEASE sanity check your outcome. To be save we reeommend editing the PDB-file untilcell dimenions to match.')

            except TypeError:
                ssio.warning_message("Warning: UnitCell lengths were not provided... This may cause issues but we're going to assume everything is OK for now...")
                                                                           
        return traj



    def __get_all_proteins(self, trajectory):
        """
        Internal function that builds a single trajectory which contains all protein
        residues. 

        Parameters
        -----------
        trajectory : mdtraj.Trajectory
            An already parsed trajectory object (i.e. checked for CAMPARI-
            relevant defects such as unitcell issues etc)

        Returns
        ----------
        ssprotein.SSProtein
            Returns a single SSProtein object
        
        """

        

        # extract full system topology
        topology = trajectory.topology

        protein_atoms = []
        
        # for each chain in this toplogy determine if the 
        # first residue is protein or not. If it's protein we parse it if 
        # not it gets skipped
        for chain in topology.chains:

            # if the first residue in the chain is protein
            # note that a formic acid cap ('FOR') is not recognized as protein
            # so we include an edgecase here for that
            if chain.residue(0).name in self.valid_residue_names:

                # intialize an empty list of atoms
                local_atoms = []

                # get every atom in this chain
                for atom in chain.atoms:
                    protein_atoms.append(atom.index)

                protein_atoms.extend(local_atoms)

        return SSProtein(trajectory.atom_slice(protein_atoms))
        

    #oxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxo
    #
    #
    def __get_proteins(self, trajectory, debug):
        """
        Internal function that takes an MDTraj trajectory and returns a list of mdtraj 
        trajectory objects corresponding to each protein in the system, ASSUMING that 
        each protein is in its own chain.
        
        The way this works is to cycle through each chain, identify if that chain
        contains protein or not, and if it does grab all the atoms in that chain
        and perform an atomslice using those atoms on the main trajectory.
        

        Parameters
        -----------
        trajectory : mdtraj.Trajectory
            An already parsed trajectory object (i.e. checked for CAMPARI-
            relevant defects such as unitcell issues etc)

        Returns
        ---------
        list
            Returns a proteinTrajectoryList - a list with one or more SSProtein
            objects in it
        
        """

        # extract full system topology
        topology = trajectory.topology



        chainAtoms = []
        
        # for each chain in this toplogy determine if the 
        # first residue is protein or not. If it's protein we parse it if 
        # not it gets skipped
        for chain in topology.chains:

            # if the first residue in the chain is protein
            # note that a formic acid cap ('FOR') is not recognized as protein
            # so we include an edgecase here for that
            if chain.residue(0).name in self.valid_residue_names:

                # intialize an empty list of atoms
                local_atoms = []

                # get every atom in this chain
                for atom in chain.atoms:
                    local_atoms.append(atom.index)

                chainAtoms.append(local_atoms)

            else:
                if debug:
                    ssio.debug_message('Skipping residue %s from %s' %(chain.residue(0).name, chain))

        # for each protein chain that we have atomic indices
        # for (hopefully all of them!) cycle through and create
        # sub-trajectories
        proteinTrajectoryList = []
        for local_chain_atoms in chainAtoms:

            # generate a trajectory composed of *JUST* the
            # $chain atoms. The PT object created now is self
            # consistent and contains an associated and fully
            # correct .topology object (NOTE this fixes a 
            # previous bug in CAMPARITraj 0.1.4)
            PT = trajectory.atom_slice(local_chain_atoms)

            # WA
            # gets the resid offset in a way that is ensures internal
            # consistency for the SSProtein object
            first_resid = PT.topology.chain(0).residue(0).index

            if first_resid != 0:
                raise SSException('After extracting a protein subtrajectory, the first resid is not 0. This may reflect a bug, or you may not be using MDTraj 1.9.5')
                
            # add that SSProtein to the ever-growing proteinTrajectory list
            proteinTrajectoryList.append(SSProtein(PT))

        if len(proteinTrajectoryList) == 0:
            ssio.warning_message('No protein chains found in the trajectory')

        return proteinTrajectoryList




    #oxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxo
    #
    #
    def __get_proteins_by_residue(self, trajectory, residue_grouping, debug):
        """
        Internal function which returns a list of mdtraj trajectory objects corresponding 
        to each protein where we *explicitly* define the residues in each protein.

        Unlike the `__get_proteins()` function, which doesn't require any manual
        input in identifying the proteins, here we provide a list of groups, where 
        each group is the set of residues associated with a protein.

        The way this works is to cycle through each group, and for each residue in 
        each group grabs all the atoms and uses these to carry out an atomslice 
        on the full trajectory. 

        Parameters
        -----------
        trajectory : mdtraj.Trajectory
            An already parsed trajectory object (i.e. checked for CAMPARI-
            relevant defects such as unitcell issues etc)

        residue_grouping : list of list of integers
            Must be a list containing one or more lists, where each internal list
            contains a set of monotonically increasing residues (which correspond
            to the full protein trajectory). In other words, each sub-list defines
            a single protein. The integer indexing here - importantly - uses the 
            CAMPARITraj internal residue indexing, meaning that indexing begins at 0 from
            the first residue in the PDB file.

        Returns
        ---------
        list
            Returns a proteinTrajectoryList - a list with one or more SSProtein
            objects in it
        

        """
        
        group_atoms = []

        # extract full system topology
        topology = trajectory.topology
        
        # for each chain in this toplogy determine if the 
        # first residue is protein or not        
        for group in residue_grouping:


            # build a string of the resids
            res_string = ''
            for r in group:
                res_string = res_string + " %i" % (int(r)) 
            
            # select atoms based on the resid string
            local_atoms = topology.select('resid %s' %(res_string))
            
            if len(local_atoms) == 0:
                ssio.warning_message('In residue group [%s ...] no residues in the trajectory were found...' %(str(group)[0:8]))

            else:
                group_atoms.append(local_atoms)

        # for each protein group that we have atomic indices
        # for cycle through and create sub-trajectories
        proteinTrajectoryList=[]
        
        # cycle through each group of atoms as was defined by the residue_grouping
        # indices

        for local_group_atoms in group_atoms:

            # generate a trajectory composed of *JUST* the
            # $chain atoms. The PT object created now is self
            # consistent and contains an associated and fully
            # correct .topology object (NOTE this fixes a 
            # previous bug in CAMPARITraj 0.1.4)
            PT = trajectory.atom_slice(local_group_atoms)       
            
            # gets the resid offset in a way that is ensures internal
            # consistency for the SSProtein object
            resid_offset = PT.topology.chain(0).residue(0).index

            if resid_offset != 0:
                raise SSException('After extracting a protein subtrajectory, the first resid is not 0. This may reflect a bug, or you may not be using MDTraj 1.9.5')
                
            proteinTrajectoryList.append(SSProtein(PT))

        if len(proteinTrajectoryList) == 0:
            ssio.warning_message('No protein chains found in the trajectory')


        return proteinTrajectoryList


    #oxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxo
    #
    #
    def get_overall_radius_of_gyration(self):
        """
        Function which returns the per-frame OVERALL radius of gyration for every 
        protein residue in the trajectory. For systems with multiple protein chains, 
        all chains are combined together. For systems with a single protein chain,
        this function offers no advantage over interacting directly with the
        SSProtein object in the .proteinTrajectoryList.

        Parameters
        -------------
        None

        Returns
        ----------
        np.ndarray 
            Returns a numpy array with per-frame instantaneous radius of gyration
        
        """

        return self.__single_protein_traj.get_radius_of_gyration()

    #oxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxo
    #
    #
    def get_overall_asphericity(self):
        """
        Function which returns the per-frame OVERALL asphericity for every
        protein residue in the trajectory. For systems with multiple protein chains, 
        all chains are combined together. For systems with a single protein chain,
        this function offers no advantage over interacting directly with the
        SSProtein object in the .proteinTrajectoryList.

        Parameters
        -------------
        None

        Returns
        ----------
        np.ndarray 
            Returns a numpy array with per-frame instantaneous asphericity
        """

        return self.__single_protein_traj.get_asphericity()

    #oxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxo
    #
    #
    def get_overall_hydrodynamic_radius(self):
        """
        Function which returns the per-frame OVERALL hydrodynamic radius for every
        protein residue in the trajectory. For systems with multiple protein chains, 
        all chains are combined together. For systems with a single protein chain,
        this function offers no advantage over interacting directly with the
        SSProtein object in the .proteinTrajectoryList.

        Parameters
        -------------
        None

        Returns
        ----------
        np.ndarray 
            Returns a numpy array with per-frame instantaneous asphericity

        """

        return self.__single_protein_traj.get_hydrodynamic_radius()


    #oxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxo
    #
    #
    def get_interchain_distance_map(self, proteinID1, proteinID2, mode='CA'):
        """        
        Function which returns two matrices with the mean and standard deviation distances
        between the residues in resID1 from proteinID1 and resID2 from proteinID2

        This computes the (full) intramolecular distance map, where the "distancemap"
        function computes the intermolecular distance map.

        Specifically, this allows the user to define two distinct chains (i.e. an "interchain"
        distance map).

        Obviously this only makes sense if your system has two separate protein objects
        defined, but in principle the output from:

            `intra_chain_distance_Map(0,0)`

        would be the same as:
        
            `proteinTrajectoryList[0].get_distance_map()`

        This is actually a useful sanity check!
        
        Parameters
        ----------

        proteinID1 : int
            The ID of the first protein of the two being considered, where the ID is the proteins position in the
            `self.proteinTrajectoryList` list. 

        proteinID2 : int
            The ID of the second protein of the two being considered, where the ID is the proteins position in the
            `self.proteinTrajectoryList` list

        mode : ['CA','COM']
            String, must be one of either 'CA' or 'COM'.
            - 'CA' = alpha carbon.
            - 'COM' = center of mass (associated withe the residue).
            Default = 'CA'.


        Returns
        -------
        tuple : tuple containing `distanceMap` and `STDMap`

        distanceMap : numpy matrix
            Is an [n x m] matrix where n and m are the number of proteinID1 residues
            and proteinID2 residues. Each position in the matrix corresponds to the
            mean distance between those two residues over the course of the simulation.

        stdMap : numpy matrix
            Is an [n x m] matrix where n and m are the number of proteinID1 residues
            and proteinID2 residues. Each position in the matrix corresponds to the
            standard devaiation associated with the distances between those two
            residues.

        
        """

        ssutils.validate_keyword_option(mode, ['CA', 'COM'], 'mode')
        
        # get SSProtein objects for the two IDs passed (could be the same)
        P1 = self.proteinTrajectoryList[proteinID1]        
        P2 = self.proteinTrajectoryList[proteinID2]

        # create the empty distance maps
        p1_residues = P1.resid_with_CA
        p2_residues = P2.resid_with_CA
        map_shape   = (len(p1_residues), len(p2_residues))
        distanceMap = np.zeros(map_shape)
        stdMap      = np.zeros(map_shape)

        for r1 in p1_residues:
            if mode == 'COM':
                COM_1 = P1.get_residue_COM(r1)
            else:
                COM_1 = P1.get_residue_COM(r1, atom_name='CA')

            p1_index = copy(r1)
            if P1.ncap:
                p1_index = r1 - 1

            for r2 in p2_residues:

                if mode == 'COM':
                    COM_2 = P2.get_residue_COM(r2)
                else:
                    COM_2 = P2.get_residue_COM(r2, atom_name='CA')

                p2_index = copy(r2)
                if P2.ncap:
                    p2_index = r2 - 1
                
                # compute distance...
                d = 10*np.sqrt(np.square(np.transpose(COM_1)[0] - np.transpose(COM_2)[0]) + np.square(np.transpose(COM_1)[1] - np.transpose(COM_2)[1])+np.square(np.transpose(COM_1)[2] - np.transpose(COM_2)[2]))

                
                distanceMap[p1_index, p2_index] =  np.mean(d, 0)
                stdMap[p1_index, p2_index]    =  np.std(d, 0)
                
        return (distanceMap, stdMap)


    #oxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxoxoxoxoxoxoxoxoxoxooxoxo
    #
    #
    def get_interchain_distance(self, proteinID1, proteinID2, R1, R2, A1='CA', A2='CA', mode='atom', periodic=False):
        """
        Function which returns the distance between two specific atoms on two residues, or between 
        two residues based on mdtraj' atomselection mode rules (discussed below). Required input are protein
        ID selectors and the resid being used. Resids should be used as would be normally used for the SSProtein
        objects associated with proteinID1 and proteinID2.

        For inter-atomic distances, atoms are selected from the passed residue and their 'name' field from the topology 
        selection language (e.g. "CA", "CB" "NH" etc). By default CA atoms are used, but one can define
        any residue of interest. We do not perform any sanity checking on the atom name - this gets 
        really hard - so have an explicit try/except block which will warn you that you've probably 
        selected an illegal atom name from the residues.

        For inter-residue distances the associated rules are defined by the 'mode' selector. By default
        mode is set to 'atom', which means the variables A1 and A2 are used (with CA as default) to define
        inter-residue distance. However, if one of the other modes are used the A1/A2 parameters are ignored
        and alternative rules for computing inter-residue distance are used. These modes are detailed below.

        Distance is returned in Angstroms.

        Parameters
        ----------

        proteinID1 : int
            Index of the first protein of interest

        proteinID2 : int
            Index of the second protein of interest

        R1 : int
            Residue index of first residue

        R2 : int
            Residue index of second residue

        A1 : str
            Atom name of the atom in R1 we're looking at. Default = 'CA'

        A2 : str
            Atom name of the atom in R2 we're looking at. Default='CA'

        mode : str
            Mode allows the user to define different modes for computing atomic distance.

            The default is 'atom' whereby a pair of atoms (A1 and A2) are provided. Other options
            are detailed below and are identical to those offered by mdtraj in compute_contacts.

            Note that if modes other than 'atom' are used the A1 and A2 options are ignored.

            + `'ca'` - same as setting `'atom'` and `A1='CA'` and `A2='CA'`, this uses the C-alpha atoms.
            + `'closest'` - closest atom associated with each of the residues, i.e. the is the point \
                            of closest approach between the two residues.
            + `'closest-heavy'` - same as `'closest'`, except only non-hydrogen atoms are considered.
            + `'sidechain'` - closest atom where that atom is in the sidechain.
            + `'sidechain-heavy'` - closest atom where that atom is in the sidechain and is heavy.

        periodic : Bool
            Flag which if distances mode is passed as anything other than 'atom' then this determines
            if the minimum image convention should be used. Note that this is only available if pdb
            crystal dimensions are provided, and in general it's better to set this to false and 
            center the molecule first. Default = False.

        Returns
        -----------
        np.array
            Returns a 1D numpy array with the distance-per-frame betwee the specified residues

        """

        # check mode keyword is valid
        allowed_modes = [ 'atom', 'ca', 'closest', 'closest-heavy', 'sidechain', 'sidechain-heavy' ]
        if mode not in allowed_modes:
            raise SSException("Provided mode keyword must be one of 'atom', 'ca', 'closest', 'closest-heavy', 'sidechain', or 'sidechain-heavy'. Provided keyword was [%s]" % (mode))

        # get SSProtein objects for the two IDs passed (could be the same)        
        try:
            P1 = self.proteinTrajectoryList[proteinID1]
            P2 = self.proteinTrajectoryList[proteinID2]

        except IndexError as e:
            raise SSException('In get_interchain_distance(): When selecting protein indices %i and %i at least one of these was out of range (indices are from 0...%i)' % (proteinID1, proteinID2, len(self.proteinTrajectoryList)-1))
            

        # next build a new trajectory that contains ONLY the two residues selected
        local_atoms1 = P1.topology.select('resid %i' % (R1))
        local_atoms2 = P2.topology.select('resid %i' % (R2))

        if len(local_atoms1) == 0:
            raise SSException("In get_interchain_distance(): When selecting resid %i from proteinID1 found no atoms" %(R1))

        if len(local_atoms2) == 0:
            raise SSException("In get_interchain_distance(): When selecting resid %i from proteinID2 found no atoms" %(R2))
            
        subtraj_p1 = P1.traj.atom_slice(local_atoms1)
        subtraj_p2 = P2.traj.atom_slice(local_atoms2)
        
        # this is now a subtrajectory which in principle contains just two residues. We can check
        # this to ensure that the trajectory has exactly 2 residues
        full_subtraj = subtraj_p1.stack(subtraj_p2)
        full_subtraj_residues = [i for i in full_subtraj.topology.residues]
        if len(full_subtraj_residues) != 2:
            raise SSException("In get_interchain_distance(): When passed in two residues (R1=%i, R2=%i) in proteins %i and %i found multiple residues (%i)...these resids could not be found " %(R1, R2, proteinID1, proteinID2, len(full_subtraj_residues)))

        
        # if we're looking at a specific pair of atoms (note we use resid 0 and 1 because we KNOW this trajectory only has 2 residues and we know R1 is 0 and R2 is 1
        if mode == 'atom':

            atom1 = full_subtraj.topology.select('resid 0 and name "%s"' % A1)
            if len(atom1) != 1:
                raise SSException("In get_interchain_distance() when selecting atom %s from residue %i in protein %i no atoms were found " % (A1, R1,  proteinID1))                

            COM_1 = md.compute_center_of_mass(full_subtraj.atom_slice(atom1))

            atom2 = full_subtraj.topology.select('resid 1 and name "%s"' % A2)
            if len(atom2) != 1:
                raise SSException("In get_interchain_distance() when selecting atom %s from residue %i in protein %i no atoms were found " % (A2, R2,  proteinID2))                

            COM_2 = md.compute_center_of_mass(full_subtraj.atom_slice(atom2))
            
            # finally compute distances
            distances = 10*np.sqrt(np.square(np.transpose(COM_1)[0] - np.transpose(COM_2)[0]) + np.square(np.transpose(COM_1)[1] - np.transpose(COM_2)[1])+np.square(np.transpose(COM_1)[2] - np.transpose(COM_2)[2]))

        else:

            # TODO: Documentation missing!
            # use the compute_contacts() function from mdtraj, multiplying by 10 because this will
            # by default give you numbers that... 
            distances = 10*md.compute_contacts(full_subtraj, [[0, 1]], scheme=mode, periodic=periodic)[0].ravel()

        return distances
      

