#!/usr/bin/python

# extract the optimized geometries from a OPT potential energy scan
# and write them to pdb files
# also, produce a listing of the structures' energies
# Olgun Guvench 2006
# Edits by Kenno Vanommeslaeghe:
# 2007: AM1 support (handy for test runs)
# 2008: do the rounding correctly for the purpose of generating pdb file names
# To do: get scan coordinate from Gaussian file by translating to python:
# grep ' Scan ' 2pdo_$i.log | awk '{print $3}'
import os
import os.path
import sys
import string

#print sys.argv

log_name = sys.argv[1]
template_name = sys.argv[2]
if sys.argv[3] == "SCF":
  efield = 4
  etype = "SCF"
  etypefield = 0
elif sys.argv[3] == "EUMP2":
  efield = 5
  etype = "EUMP2"
  etypefield = 3
elif sys.argv[3] == "Energy":
  efield = 1
  etype = "Energy="
  etypefield = 0
else:
  print "argv[3] must be either SCF, EUMP2 or Energy"
  quit

dihedrala = sys.argv[4]

in_file = open( template_name, "r")
natom = 0
template = []
for line in in_file.readlines():
  field = string.split( string.strip( line ) )
  if field[0] == "ATOM":
    natom = natom + 1
    template.append( line )
  elif field[0] == "TER":
    template.append( line )
  elif field[0] == "END":
    template.append( line )
in_file.close()

if os.access( "%s_pdb" % ( log_name ), os.F_OK ):
  for root, dirs, files in os.walk( "%s_pdb" % ( log_name ), \
   topdown=False):
    for name in files:
      os.remove(os.path.join(root, name))
    for name in dirs:
      os.rmdir(os.path.join(root, name))
  os.rmdir( "%s_pdb" % ( log_name ) )
os.mkdir( "%s_pdb" % ( log_name ) )

in_file = open( log_name, "r")

e_file = open( "e_file.txt" , "w" )
s_file = open( "s_file.txt" , "w" )
c_file = open( "c_file.txt" , "w" )
line = in_file.readline()
while line:
# print string.rstrip( line )
  field = string.split( string.strip( line ) )
  if len( field ) > 1:
    if len( field ) > etypefield and field[etypefield] == etype:
      emp2 = float( field[efield].replace("D", "E") ) * 627.5095
    elif len( field ) > 1 and field[0] == "Input" and field[1] == "orientation:":
      iatom = 1
      while iatom <= 4:
#       print in_file.readline()
        in_file.readline()
        iatom = iatom + 1

      iatom = 0
      x = []
      y = []
      z = []
      while iatom < natom:
        field = string.split( string.strip( in_file.readline() ) )
#       print field
        x.append( float( field[3] ) )
        y.append( float( field[4] ) )
        z.append( float( field[5] ) )
        iatom = iatom + 1
    elif field[0] == "Optimization" and ( field[1][0:9] == "completed" or field[1][0:7] == "stopped" ):
      line = in_file.readline()
      while line: 
        field = string.split( string.strip( line ) )
        if len( field ) > 4 and field[2] == dihedrala and field[4] == "-DE/DX":
          dihedral = float( field[3] )
          break
        line = in_file.readline()
#Bugfix: use %g with round(x) instead of %i with int(x+0.1)
      dihedraltag="%g" % (0.1*round(dihedral*10))
      out_file = \
       open( "%s_pdb/%s.%s.pdb" % ( log_name, log_name, dihedraltag ), "w")
      out_file.write( "REMARK %s = %f, %s %f kcal/mol\n" \
       % ( dihedrala, dihedral, etype, emp2 ) )
      e_file.write( "%s\n" % ( emp2 ) )
      s_file.write( "%s\n" % ( os.path.abspath( "%s_pdb/%s.%s.pdb" % ( log_name, log_name, dihedraltag ) ) ) )
      iatom = 0
      while iatom < natom:
        c_file.write( "%.5f %.5f %.5f\n" \
        % ( x[iatom], y[iatom], z[iatom] ) )
        out_file.write( "%s%8.3f%8.3f%8.3f%s" \
        % ( template[iatom][0:30], x[iatom], y[iatom], z[iatom], template[iatom][54:] ) )
        iatom = iatom + 1
      out_file.write( "%s" % ( template[iatom] ) )
      out_file.write( "%s" % ( template[iatom+1] ) )
      out_file.close()

  line = in_file.readline()


