from MDAnalysis import *
import MDAnalysis.analysis
import numpy
import sys



def numOfNeighbours(universe,selection1,selection2,radius):
    """Number of neighbours"""

    # sum of neighbours of all residues
    sumneighbours = 0

    # create atomgroups
    grp1 = universe.selectAtoms(selection1)
    grp2 = universe.selectAtoms(selection2)
    grp_both = grp1+grp2

    # loop through all residues of grp1
    num_of_residues = grp1.numberOfResidues()
    atoms_per_residue = grp1.numberOfAtoms() / num_of_residues
    for ritr in range(0, num_of_residues-1):

        # range of atom indices of this residue
        first_atom = ritr * atoms_per_residue + 1
        last_atom = (ritr+1) * atoms_per_residue

        # indice of last atom in this residue
        last_atom = first_atom + atoms_per_residue - 1

        # select all atoms close to any atom of this residue
        sel_around = 'around ' + str(radius) + ' (bynum ' + str(first_atom) + ':' + str(last_atom) + ')'
        grp_around = grp_both.selectAtoms(sel_around)
        if (grp_around.numberOfAtoms() < 1):
            continue

        # select atoms contained in both grp2 and grp_around
        grp_neighbours = grp_around.selectAtoms(selection2)
        if (grp_neighbours.numberOfAtoms() < 1):
            continue

        # sum residue-residue contacs found for this residue
        sumneighbours = sumneighbours + grp_neighbours.numberOfResidues()

    return sumneighbours




def contacts(coord,traj,rad,select1,select2,first_ts,last_ts):

    universe = Universe(coord,traj)
    outfilename = 'contacts_' + str(first_ts) +'-' + str(last_ts) + '.xvg'

    print 'rad: ' + str(rad)
    print 'first_ts: ' + str(first_ts)
    print 'last_ts: ' + str(last_ts)
    print 'select1: ' + select1
    print 'select2: ' + select2
    print 'outfilename: ' + outfilename

    for ts in universe.trajectory:
        if (ts.time < first_ts or ts.time >= last_ts):
            print 'skip ' + str(ts.time)
            continue
    
        result = str(numOfNeighbours(universe,select1,select2,rad))
        timestep = str(ts.time)
        print timestep + ' ' + result
        
        # write file
        newline = '{:>12} {:>12}\n'.format(timestep, result)
        outfile = open(outfilename, 'a+')
        outfile.write(newline)
        outfile.close()