#!@PYTHON@
#
# Removes random atoms from a TREMOLO data file

import sys, random, math, re

# check arguments
if len(sys.argv) < 5:
  print "Usage: "+sys.argv[0]+" <datafile> <NoDefects> <OffsetElement> <OffsetXYZ> <BondDistance>"
  sys.exit(1)

input = open(sys.argv[1], "r")
output = open(sys.argv[1]+".defect", "w")
defectcount=int(sys.argv[2])
OffsetElement=int(sys.argv[3])
OffsetXYZ=int(sys.argv[4])
DISTANCE=float(sys.argv[5])
THRESHOLD=.4
CUTOFF=DISTANCE+THRESHOLD

# 1. parse lines into array
lines=[]
comments=[]
nr=0
for line in input:
  if "#" in line:
    comments.append(line)
    continue
  lines.append(line.split())
nr=len(lines)
input.close()
print "I have scanned "+str(nr)+" lines in "+sys.argv[1]+" with "+str(len(comments))+" comments."

# 2a. parse atoms coordinates into list and find bounding box
max=[ 0.,  0., 0. ] 
min=[ 0.,  0., 0. ] 
AllAtoms=[]

for n in range(len(lines)):
  atom=lines[n]
    #print"reading atom"
  for i in range(3):
    #print float(atom[i+1])
    if float(atom[i+OffsetXYZ]) < min[i]:
      min[i] = float(atom[i+1])
    if float(atom[i+OffsetXYZ]) > max[i]:
      max[i] = float(atom[i+OffsetXYZ])
  AllAtoms.append([n, atom[OffsetElement], float(atom[OffsetXYZ]), float(atom[OffsetXYZ+1]), float(atom[OffsetXYZ+2])])
print "Found Box bounds [%f,%f]x[%f,%f]x[%f,%f]." % (min[0],max[0],min[1],max[1],min[2],max[2])

# 2b. get number of linked cells
cells_x=int(math.ceil((max[0]-min[0])/CUTOFF))
cells_y=int(math.ceil((max[1]-min[1])/CUTOFF))
cells_z=int(math.ceil((max[2]-min[2])/CUTOFF))
print "Number of cells in each axis direction (%f,%f,%f)." % (cells_x, cells_y, cells_z)

# 2c. create 3 dim array to fit in box, each cell has size CUTOFF Ang
cell=[]
for i in range(cells_x):
  cell.append([])
  for j in range(cells_y):
    cell[i].append([])
    for k in range(cells_z):
      cell[i][j].append([0])

# 2d. parse every atom into appropriate cell
n=0
for atom in AllAtoms:
  x=int(math.floor((atom[2]-min[0])/CUTOFF))
  y=int(math.floor((atom[3]-min[1])/CUTOFF))
  z=int(math.floor((atom[4]-min[2])/CUTOFF))
  cell[x][y][z][0] +=1
  cell[x][y][z].append(atom[0])		# append index to AllAtoms list
  n=n+1
print str(n)+" of "+str(len(AllAtoms))+" atoms were incarcerated."

# 2e. go through every cell, check neighbours
atom1=[0,0.,0.,0.]
atom2=[0,0.,0.,0.]
v=0
NumBonds=0
Bonds=[]
for i in range(cells_x):
  for j in range(cells_y):
    for k in range(cells_z):
      v=cell[i][j][k][0]
      #go through every atom in cell
      for l in range(1, v+1):
        atom1=cell[i][j][k][l]          
        # go through cell and all lower neighbours
        for e in range(i-1,i+2):
          #if on border continue periodic
          if e>=cells_x or e<0:
            #e=e-cells_x
            continue
          for r in range(j-1,j+2):
            #if on border continue periodic
            if r>=cells_y or r<0:
              #r=r-cells_y
              continue
            for t in range(k-1,k+2):
              #if on border continue periodic
              if t>=cells_z or t<0:
                #t=t-cells_z
                continue
              #efficiency in linked cell: ommit four boxes in middle layer
              #if (i==e and r==j+1):
              #  continue
              #if (i==e and r==j+1 and t==k+1):
              #  continue
              w=cell[e][r][t][0]
              #go through all atoms in cell
              for m in range(1, w+1):
                atom2=cell[e][r][t][m]
                #when in same cell: ommit identical particles
                if (AllAtoms[atom1][0]>=AllAtoms[atom2][0]):
                  continue
                dist=0
                for h in range(3):
                  dist += (AllAtoms[atom1][2+h] - AllAtoms[atom2][2+h])*(AllAtoms[atom1][2+h] - AllAtoms[atom2][2+h]) 
                dist=math.sqrt(dist)

                # add bond
                if (dist >= (DISTANCE-THRESHOLD)) and (dist <= (DISTANCE+THRESHOLD)):
                  Bonds.append([atom1, atom2])
                  NumBonds+=1
print str(NumBonds)+" bonds with partners distanced within ["+str(DISTANCE-THRESHOLD)+","+str(DISTANCE+THRESHOLD)+"] angstroem were found."

# 3. check whether we have sufficient bonds and choose desired number
if NumBonds < defectcount:
  print "Not enough bonds for the desired number of defects."
  sys.exit(1)

# 4. flip random but unique bonds by 90 degrees
vector1=[0., 0., 0.]
vector2=[0., 0., 0.]
for n in range(defectcount):
  # 4a. determine random bond and its bonding partners
  atom3=-1
  while atom3==-1: # pick a unique bond
    i=random.randrange(0,NumBonds)
    atom1=Bonds[i][0]
    atom2=Bonds[i][1]
    atom3=-1
    if (atom1 == -1) or (atom2 == -1):
      continue
    # 4b. obtain a third atom to describe the plane fully
    for j in range(NumBonds):
      if ((Bonds[j][0] != atom2) and (Bonds[j][1] == atom1)) or ((Bonds[j][0] != atom1) and (Bonds[j][1] == atom2)):
        atom3=Bonds[j][0]
        break
      if ((Bonds[j][1] != atom2) and (Bonds[j][0] == atom1)) or ((Bonds[j][1] != atom1) and (Bonds[j][0] == atom2)):
        atom3=Bonds[j][1]
        break
    if (atom3 == -1):
      continue
  print "Flipping bond "+str(i)+"."
  Bonds[i][0]=-1
  Bonds[i][1]=-1
  # 4c. determine in-plane vectors and orthonormalize
  for i in range(3):
    vector1[i]=AllAtoms[atom1][2+i]
    vector2[i]=AllAtoms[atom1][2+i]
  for i in range(3):
    vector1[i]-=AllAtoms[atom2][2+i]
    vector2[i]-=AllAtoms[atom3][2+i]
  BondNorm=0
  for i in range(3):
    BondNorm+=vector1[i]*vector1[i]
  BondNorm=math.sqrt(BondNorm)
  for i in range(3):
    vector1[i]/=BondNorm
  norm=0
  for i in range(3):
    norm+=vector1[i]*vector2[i]
  for i in range(3):
    vector2[i]-=norm*vector1[i]
  norm=0
  for i in range(3):
    norm+=vector2[i]*vector2[i]
  norm=math.sqrt(norm)
  for i in range(3):
    vector2[i]/=norm
  vector4=[0., 0, 0.] 
  for i in range(3):
    vector4[i]=(AllAtoms[atom1][2+i]+AllAtoms[atom2][2+i])/2.
    vector2[i]*=BondNorm/2.

  # 4b. rotate the bonding nuclei and store in array lines!
  for i in range(3):
		lines[atom1][OffsetXYZ+i]="%e" % (vector4[i]+vector2[i])
		lines[atom2][OffsetXYZ+i]="%e" % (vector4[i]-vector2[i])

# 5. store AllAtoms to file
for line in comments:
  output.write("%s" % (line))
for line in lines:
  for n in range(len(line)):
    output.write("%s\t" % (line[n]))
  output.write("\n")
output.close()  

# 6. exit
print "All done."

sys.exit(0)