#! /usr/bin/env python
# Copyright (c) 2008 Robert L. Campbell (rlc1@queensu.ca)
#
import sys,math,numpy


def coord_input(coord_file_name):
  coord_infile = coord_file_name.replace('.xyzrn','')
  coord_infile = open(coord_file_name + '.xyzrn')
  atom_list = []
  for line in coord_infile.readlines():
    x,y,z,rad,flag,atom_name = line.strip().split()
    atom_list.append(((float(x),float(y),float(z)),float(rad),flag,atom_name))
  return atom_list

def vertex_input(vertex_file_name,file_out):
  vertex_file_name = vertex_file_name.replace('.triang.vert','')
  vertex_file_name = vertex_file_name.replace('.vert','')
  vert_infile = open(vertex_file_name + '.vert')
# strip off header lines
  vert_line1 = vert_infile.readline()
  vert_line2 = vert_infile.readline()
  vert_line3 = vert_infile.readline()
  file_out.write(vert_line1)
  file_out.write(vert_line2)
  file_out.write(vert_line3)
  vertices = []
  for line in vert_infile.readlines():
    data = line.strip().split()
    vertex = [float(x) for x in data[0:3]]
    normal = [float(x) for x in data[3:6]]
    face_number = int(data[6])
    atom_index = int(data[7])
    toric_flag = int(data[8])
    atom_name = data[9]
    vertices.append((vertex,normal,face_number,atom_index,toric_flag,atom_name))
  return vertices

def main():

  weight = .5
# read first vertex files
  file_name1 = sys.argv[1]
  file_name2 = sys.argv[2]
  file_out1 = open(file_name1 + '_' + file_name2 + '.vert','w')
  file_out2 = open(file_name2 + '_' + file_name1 + '.vert','w')

  atoms1 = coord_input(file_name1)
  vertices1 = vertex_input(file_name1,file_out1)

  atoms2 = coord_input(file_name2)
  vertices2 = vertex_input(file_name2,file_out2)

# create interatomic distance squared array so inter-atomic distance 
# calculation is not done once for every vertex pair, but only for every atom pair
  #dist_array = make_distance_array(atoms1,atoms2)
  dist_array = numpy.zeros((len(atoms1),len(atoms2)),dtype=numpy.float)
  for i in xrange(len(atoms1)):
    for j in xrange(len(atoms2)):
      dist_array[i][j] = (atoms1[i][0][0]-atoms2[j][0][0])**2 + (atoms1[i][0][1]-atoms2[j][0][1])**2 + (atoms1[i][0][2]-atoms2[j][0][2])**2

# begin loop over vertices from surface 1 to surface 2
  s1_2_array = []
  for i in xrange(len(vertices1)):
    min_dist1sq = 99999.
    vertex1 = vertices1[i][0]
    normal1 = vertices1[i][1]
    atom_index1 = vertices1[i][3]
    closest_index2 = -1
    for j in xrange(len(vertices2)):
      vertex2 = vertices2[j][0]
      normal2 = vertices2[j][1]
      atom_index2 = vertices2[j][3]
# check atom distance from dist_array first 
# (dist_array is zero based, atom numbers in vertex file start at 1)
      if dist_array[atom_index1-1][atom_index2-1] < 64:
        distsq = (vertex1[0]-vertex2[0])**2 + (vertex1[1]-vertex2[1])**2 + (vertex1[2]-vertex2[2])**2
        if distsq < min_dist1sq:
          closest_index2 = j
          min_dist1sq = distsq

    if closest_index2 != -1:
      a1 = normal1
      a2 = vertices2[closest_index2][1]
      dot_prod = (a1[0]*-a2[0]+a1[1]*-a2[1]+a1[2]*-a2[2])
      s1_2 = dot_prod * math.exp(-weight * min_dist1sq)
      s1_2_array.append(s1_2)
    else:
      s1_2 = -1

    file_out1.write("%9.3f %9.3f %9.3f %9.3f %9.3f %9.3f %7d %7d %2d %s %9.3f %9.3f\n" % (
           vertices1[i][0][0],
           vertices1[i][0][1],
           vertices1[i][0][2],
           vertices1[i][1][0],
           vertices1[i][1][1],
           vertices1[i][1][2],
           vertices1[i][2],
           vertices1[i][3],
           vertices1[i][4],
           vertices1[i][5],
           min_dist1sq, s1_2 ))

# begin loop over vertices from surface 2 to surface 1
  s2_1_array = []
  for i in xrange(len(vertices2)):
    min_dist2sq = 99999.
    vertex2 = vertices2[i][0]
    normal2 = vertices2[i][1]
    atom_index2 = vertices2[i][3]
    closest_index1 = -1
    for j in xrange(len(vertices1)):
      vertex1 = vertices1[j][0]
      normal1 = vertices1[j][1]
      atom_index1 = vertices1[j][3]
# check atom distance from dist_array first 
# (dist_array is zero based, atom numbers in vertex file start at 1)
      if dist_array[atom_index1-1][atom_index2-1] < 64:
        distsq = (vertex1[0]-vertex2[0])**2 + (vertex1[1]-vertex2[1])**2 + (vertex1[2]-vertex2[2])**2
        if distsq < min_dist2sq:
          closest_index1 = j
          min_dist2sq = distsq

    if closest_index1 != -1:
      a1 = normal2
      a2 = vertices1[closest_index1][1]
      dot_prod = (a1[0]*-a2[0]+a1[1]*-a2[1]+a1[2]*-a2[2])
      s2_1 = dot_prod * math.exp(-weight * min_dist2sq)
      s2_1_array.append(s2_1)
    else:
      s2_1 = -1

    file_out2.write("%9.3f %9.3f %9.3f %9.3f %9.3f %9.3f %7d %7d %2d %s %9.3f %9.3f\n" % (
           vertices2[i][0][0],
           vertices2[i][0][1],
           vertices2[i][0][2],
           vertices2[i][1][0],
           vertices2[i][1][1],
           vertices2[i][1][2],
           vertices2[i][2],
           vertices2[i][3],
           vertices2[i][4],
           vertices2[i][5],
           min_dist2sq, s2_1 ))

  s1_2_median = numpy.median(s1_2_array)
  s2_1_median = numpy.median(s2_1_array)
  sys.stdout.write("Median Sc (complementarity) value for surface 1 to surface 2 is: %6.3f\n" % s1_2_median)
  sys.stdout.write("Median Sc (complementarity) value for surface 2 to surface 1 is: %6.3f\n" % s2_1_median)
  sys.stdout.write("Average median Sc (complementarity) value for surfaces 1 and 2 are: %6.3f\n" % ((s1_2_median+s2_1_median)/2))

if __name__ == "__main__":
  main()