#!@PYTHON@
#
# Cuts a specific block from a given PDB file

import sys, string, re, math
wrerr = sys.stderr.write
wrout = sys.stdout.write

def SKP(a,b):
  """ Calculates scalar product of two 3dim vectors
"""
  result=0
  for i in range(3):
    result+=a[i]*b[i]
  return(result)

# check for parameters
if len(sys.argv) < 7:
  wrerr("Usage: "+sys.argv[0]+" <PDB file> <offset to id> <offset to x> <offsetSeqID> <offset2neighbours> <offset2improper>\n")
  sys.exit(1)

input = open(sys.argv[1], "r")
output = open(sys.argv[1]+".blk.pdb", "w")
offsetID=int(sys.argv[2])
offset=int(sys.argv[3])
offsetSeqID=int(sys.argv[4])
offsetNeighbours=int(sys.argv[5])
offsetImproper=int(sys.argv[6])

# 1. determine box boundary vectors
a=[ 115., 0., 30. ]	# x
b=[ 150., 0., 30. ] # offset
c=[ 150., 64., 30. ] # y
d=[ 150., 0., 66. ] # z
origin=[ 115., 0., 30. ]
line1=[0., 0., 0.]
line2=[0., 0., 0.]
line3=[0., 0., 0.]
norma=0
normb=0
normc=0
for i in range(3):
  line1[i]=a[i]-b[i]
  line2[i]=c[i]-b[i]
  line3[i]=d[i]-b[i]
  norma+=line1[i]*line1[i]
  normb+=line2[i]*line2[i]
  normc+=line3[i]*line3[i]

# 2a. prepare mapping by removing non-trippled TIP3s
ids={}
for line in input:
  if "#" in line:		# skip comments
    continue
  entries=line.split()
  e=[ float(entries[offset]), float(entries[offset+1]), float(entries[offset+2]) ]
  for i in range(3):
    e[i]-=b[i]
  skp1 = SKP(line1,e)/norma
  skp2 = SKP(line2,e)/normb
  skp3 = SKP(line3,e)/normc
  if ((skp1 >=0) and (skp1 < 1)) and ((skp2 >=0) and (skp2 <1)) and ((skp3 >=0) and (skp3 <1)):
    if "TIP3" in line:	# note down the number of times this SeqID appears, should be 3 for TIP3 clusters (h2o)
      i=int(entries[offsetSeqID])
      if ids.has_key(i):
        ids[i]=ids[i]+1
      else:
        ids[i]=1
      #wrout("for key %d: %d present\n" % (i, ids[i]))

# 2b. assign mapping between old and new ids
input.seek(0)
id=0
idmap={}
for line in input:
  if "#" in line:		# skip comments
    continue
  entries=line.split()
  e=[ float(entries[offset]), float(entries[offset+1]), float(entries[offset+2]) ]
  for i in range(3):
    e[i]-=b[i]
  skp1 = SKP(line1,e)/norma
  skp2 = SKP(line2,e)/normb
  skp3 = SKP(line3,e)/normc
  if ((skp1 >=0) and (skp1 < 1)) and ((skp2 >=0) and (skp2 <1)) and ((skp3 >=0) and (skp3 <1)):
    if (("TIP3" in line) and (ids[int(entries[offsetSeqID])] != 3)):	# drop TIP3 that are not present as triple
      continue
    id+=1
    idmap[int(entries[offsetID])]=id	# note down old id to new id mapping

wrout("The following ids will be removed due to NON-Triple-TIP3 criterium: ")
for key in ids.keys():
  if ids[key] < 3:
    wrout("(%d = %d) " % (key, ids[key]))
wrout("\n")


# 3. write out the block with changed ids and neighbours and impropers
input.seek(0)
for line in input:
  if "#" in line:		# skip comments
    continue
  entries=line.split()
  e=[ float(entries[offset]), float(entries[offset+1]), float(entries[offset+2]) ]
  for i in range(3):
    e[i]-=b[i]
  skp1 = SKP(line1,e)/norma
  skp2 = SKP(line2,e)/normb
  skp3 = SKP(line3,e)/normc
  if ((skp1 >=0) and (skp1 < 1)) and ((skp2 >=0) and (skp2 <1)) and ((skp3 >=0) and (skp3 <1)): # check whether coordinate is inside desired block
    if (("TIP3" in line) and (ids[int(entries[offsetSeqID])] != 3)):	# drop TIP3 that are not present as triple
      continue
    i=0
    while (i < len(entries)):
      if ((i == offsetID)): # change id to new id
        output.write("%d\t" % (idmap[int(entries[i])]))
        i+=1
        continue
      if ((i >= offset) and (i < offset+3)): # shift block to start at origin
        output.write("%f\t" % (float(entries[i])-origin[i-offset]))
        i+=1
        continue
      if ((i >= offsetNeighbours) and (i < offsetNeighbours+4)): # change neighbour ids to new ods
        if idmap.has_key(int(entries[i])):
          output.write("%d\t" % ( idmap[int(entries[i])] ))
        else:
          output.write("0\t")
        i+=1
        continue
      if ((i == offsetImproper)):  # change improper ids to new ids
        if entries[i] != "-":		# if it just contains "-" (i.e. no impropers), write "-"
          neighbours=entries[i].split(',') # there may be several improper terms, separated by ","
          for k in range(len(neighbours)):
            neighbour=neighbours[k].split('-') # the improper neighbours are separated by "-"
            invalid=0
            for j in range(len(neighbour)):  # first check if each improper neighbour is present in the cut-out block
              if not idmap.has_key(int(neighbour[j])):
                invalid=1
                break
            if invalid == 0: # if all are present, map old ids to new ids
              if (k != 0):
                output.write(",")
              for j in range(len(neighbour)):
                try:
                  if idmap.has_key(int(neighbour[j])):
                    if j != 0:
                      output.write("-%d" % ( idmap[int(neighbour[j])] ))
                    else:
                      output.write("%d" % ( idmap[int(neighbour[j])] ))
                except ValueError:
                  wrout("%s is not a valid neighbour\n" % (neighbour[j]))
              # next improper term for this atom
            else:
              if (k == len(neighbours)-1):
                output.write("-")
          # next improper neighbour
          i+=1
          continue
      output.write("%s\t" % (entries[i]))
      i+=1
    output.write("\n")
    id+=1