#
# written  by Nikolaj Otte at the MPI in Muelheim
#
# This script writes a chemshell .trj trajectory into a
#binary .dcd file that can easily be viewed using either 
# gOpenmol, vmd, or insight. 
#note that it assumes to find the source trajectory as "dynamics.trj",
#the result will be "dynamics.dcd". 
#since the information is not present in the .trj, this script
#cannot properly set the stepsize nor the first step variables. dummies
#take their place.

# usage: python trj2dcd.py

import sys
import string
import re
import os
import binascii
import struct
import array
    
input = open("dynamics.trj")
output = open("dynamics.dcd",'w')

blockident = re.compile("^block.+")
coordident = re.compile("coord")
isNotCoor = 0

#get all the info from the .trj file
data = []
atomSetx = []
atomSety = []
atomSetz = []

atomx = []
atomy = []
atomz = []

#number of coordinate files
numset = 0

#starting timestep number NOTE: just a housenumber.
#set to different INTEGER value if desired
startstep = 1000

#timesteps between coordinate steps NOTE: just a housenumber.
#set to different INTEGER value if desired.
interv = 1000

#number of atoms
numat =  0
dataraw= []
input.readline()
input.readline()

for s in range(1000000):
    dataraw.append(input.readline())
    if(dataraw[s] == ''):
        break
    data.append(string.split(dataraw[s]))
    endblock = blockident.search(data[s][0])
    if endblock:
        numset = numset + 1
        if not coordident.search(dataraw[s]):
          isNotCoor = 1 
          print ("NExt are not coors!")
        break
    else:
        atomx.append(float(data[s][1])*0.529)
        atomy.append(float(data[s][2])*0.529)
        atomz.append(float(data[s][3])*0.529)
        numat = numat + 1
    

#write dcd header
output.write(struct.pack('>i',84))
output.write(struct.pack('>ssss','C','O','R','D'))
output.write(struct.pack('>i',numset))
output.write(struct.pack('>i',startstep))
output.write(struct.pack('>i',interv))
output.write(struct.pack('>i',1000))


for s in range(5):
    output.write(struct.pack('>i',0))

#hypo time step
output.write(struct.pack('>f',0.01))

output.write(struct.pack('>i',0))
output.write(struct.pack('>i',0))


for s in range(7):
    output.write(struct.pack('>i',0))


output.write(struct.pack('>i',27))
output.write(struct.pack('>i',84))
output.write(struct.pack('>i',164))
output.write(struct.pack('>i',2))

output.write(struct.pack('>ssss','f','i','l','e'))
output.write(struct.pack('>sss',' ','i','s'))
output.write(struct.pack('>ss',' ','a'))
output.write(struct.pack('>sssssssssss',' ','t','r','a','n','s','l','a','t','e','d'))
output.write(struct.pack('>ssssssssss',' ','c','h','e','m','s','h','e','l','l'))
output.write(struct.pack('>sssssssssss',' ','t','r','a','j','e','c','t','o','r','y'))
#41 bytes written
#write the 164-41 bytes
for s in range(119):
    output.write(struct.pack('>s',' '))


output.write(struct.pack('>i',164))
output.write(struct.pack('>i',4))
output.write(struct.pack('>i',numat))
output.write(struct.pack('>i',4))


output.write(struct.pack('>i',4*numat))
for i in range(numat):
    output.write(struct.pack('>f',atomx[i]))
output.write(struct.pack('>i',4*numat))
output.write(struct.pack('>i',4*numat))
for i in range(numat):
    output.write(struct.pack('>f',atomy[i]))
output.write(struct.pack('>i',4*numat))
output.write(struct.pack('>i',4*numat))
for i in range(numat):
    output.write(struct.pack('>f',atomz[i]))
output.write(struct.pack('>i',4*numat))
            
dataraw = []
data = []
atomx = []
atomy = []
atomz = []

iter = 1
value = 1
for x in range(10000):
    if isNotCoor == 0:
      for s in range(numat):
          dut = input.readline()
          if dut == '':
              value = 0
              break
          else:
              data.append(string.split(dut))
              atomx.append(float(data[s][1])*0.529)
              atomy.append(float(data[s][2])*0.529)
              atomz.append(float(data[s][3])*0.529)
      print ("iteration: %i is Coordinates" % (iter))
      if value == 0:
          break
      else:
          output.write(struct.pack('>i',4*numat))
          for i in range(numat):
              output.write(struct.pack('>f',atomx[i]))
          output.write(struct.pack('>i',4*numat))
          output.write(struct.pack('>i',4*numat))
          for i in range(numat):
              output.write(struct.pack('>f',atomy[i]))
          output.write(struct.pack('>i',4*numat))
          output.write(struct.pack('>i',4*numat))
          for i in range(numat):
              output.write(struct.pack('>f',atomz[i]))
          output.write(struct.pack('>i',4*numat))
          numset = numset + 1
          atomx = []
          atomy = []
          atomz = []
          data = []
    else: 
      print ("iteration: %i is NOT Coordinates" % (iter))         
      for s in range(numat):
          dut = input.readline()
          if dut == '':
              value = 0
              break
      if value == 0:
          break
    a = input.readline()
    if coordident.search(a):
      isNotCoor = 0   
    else:
      isNotCoor = 1    
    iter = iter + 1

output.seek(0)
output.write(struct.pack('>i',84))
output.write(struct.pack('>ssss','C','O','R','D'))
output.write(struct.pack('>i',numset))