#!/usr/bin/env python
# -*- coding: iso-8859-1 -*-

#***********************************************************#
# aoforce2hc.py, v1.0, Mikael Johansson 18.2.2003           #
#                                                           #
# This program reads a Turbomole aoforce output file and    #
# creates a HyperChem script file which then can be used by #
# HyperChem for animating the vibrations.For this you first #
# need to read in a HC .hin-file of the molecule            #
#                                                           #
# Version history:                                          #
#   v1.0: First release                                     #
#                                                           #
# Check for new versions and contact information at:        #
# http://www.iki.fi/~mpjohans/                              #
#***********************************************************#

import sys, string
split=string.split

VERSION='v1.0'
DATE='18.2.2003'
AAU=0.529177208319 # Ångström per a.u. (CODATA 1998)
ATOMS=[
# 1    2    3    4    5    6    7    8    9   10   11   12   13   14   15   16   17   18  <-IUPAC column nr.
'q' , #                                                                                     IUPAC row nr.
'h' ,                                                                                'he', # 1
'li','be',                                                  'b' ,'c' ,'n' ,'o' ,'f' ,'ne', # 2
'na','mg',                                                  'al','si','p' ,'s' ,'cl','ar', # 3
'k' ,'ca','sc','ti','v' ,'cr','mn','fe','co','ni','cu','zn','ga','ge','as','se','br','kr', # 4
'rb','sr','y' ,'zr','nb','mo','tc','ru','rh','pd','ag','cd','in','sn','sb','te','i' ,'xe', # 5
'cs','ba',                                                                                 # 6
          'la','ce','pr','nd','pm','sm','eu','gd','tb','dy','ho','er','tm','yb','ly',      # lanthanides
               'hf','ta','w' ,'re','os','ir','pt','au','hg','tl','pb','bi','po','at','rn', # 6 cont.
'fr','ra',                                                                                 # 7
          'ac','th','pa','u' ,'np','pu','am','cm','bk','cf','es','fm','md','no','lr',      # actinides
               'rf','db','sg','bh','hs','mt']                                              # 7 cont.

class AtomClass:
    def __init__(self):
        self.type=None    # atom type, ex. 'h'
        self.number=None  # atom type number, ex. 1
        self.x=None       # x-coordinate
        self.y=None         
        self.z=None

class SpectraClass:
    def __init__(self):
        self.sym=None       # symmetry (not read in at this point!)
        self.freq=None      # frequency
        self.intensity=None # IR-intensity
        self.redmass=None   # reduced mass in a.u.
        self.x=[]           # list of x-displacements for all atoms in molecule
        self.y=[]
        self.z=[]
         

def usage():
    print "\naoforce2hc "+VERSION+" by Mikael Johansson "+DATE+"\n"
    print "Usage: "+sys.argv[0]+" [-0] [-s <scale>] <inputfile> <outputfile>"
    print "                 -0: if specified keep zero-frequencies"
    print "         -s <scale>: scale displacements by <scale> (default 1.0)"
    print "        <inputfile>: aoforce output file"
    print "       <outputfile>: target HyperChem script file (*.scr recommended)"
 

def init():
    args=sys.argv
        
    zeros=0
    scale=1.0
    eline=0
    while len(args)>3:
        eline=1
        if args[1]=='-0':
            zeros=1
            print "-0 specified: keeping zero-freqs."
            del args[1]
            continue
        elif args[1]=='-s':
            scale=float(args[2])
            print "-s specified: scaling factor",scale
            del args[1:3]
            continue
        else:
            sys.stderr.write("Error: Unknown option: "+args[1]+"\n")
            usage()
            sys.exit(1)
    
    if eline:
        print
    if len(args)<>3:
        sys.stderr.write("Error: No output file specified\n")
        usage()
        sys.exit(1)
    
    infile=open(args[len(args)-2],'r')
    outfile=open(args[len(args)-1],'w')

    return infile, outfile, zeros, scale
    

#*******************
# def GetAoforceAtoms()
#
# will return list 'atoms' of following structure:
# atoms[i]=[x-coord, y-coord, z-coord, atomtype(char), atomnumber]
#*********************************
def GetAoforceAtoms(infile):

    found=0
    line=infile.readline()
    while line:
        if string.find(line,"atomic coordinates")<>-1:
            found=1
            break
        line=infile.readline()

    if not(found):
        sys.stderr.write("Error: unable to find atomic coordinates\n")
        sys.exit(2)
    

    column=split(infile.readline())
    atoms=[]


    AN=0
    while len(column)>=4:
        
        column[3]=string.lower(column[3])
        if ATOMS.count(column[3])==0:
            sys.stderr.write("Warning: Unknown atomtype found: "+column[3]+", replaced with a hydrogen\n")
            column[3]='h'
        atoms.append(AtomClass())
        atoms[AN].type=column[3]
        atoms[AN].number=ATOMS.index(column[3])
        atoms[AN].x=column[0]
        atoms[AN].y=column[1]
        atoms[AN].z=column[2]
        column=split(infile.readline())
        AN=AN+1

    return atoms


#****************
# def GetAoforceLesSpectra()
#
# Extracts IR-spectrum data from 'infile', expects aoforce "$les" output format
# Will automatically set intensities to zero, as they are not calculated!
#*********
def GetAoforceLesSpectra(infile):

    found=0
    found=0
    line=infile.readline()
    while line:
        if string.find(line,"VECTORS")<>-1:
            found=1
            break
        line=infile.readline()

    if not(found):
        sys.stderr.write("Error: No frequencies found\n")
        sys.exit(2)

    
    spectra=[]
    
    totf=0

    # First, read in frequencies and displacements from <infile> (aoforce.out)

    while 1:
        
        found=0
        line=infile.readline()
        while line:
            if string.find(line,"eigenvalue")<>-1:
                found=1
                break
            line=infile.readline()
        if not(found): # (hopefully) all modes extracted
            break
            

        column=split(line)
        spectra.append(SpectraClass())

        spectra[totf].sym="-"   #   Not extracted yet!
        
        freq=float(column[len(column)-2]) # frequency in this col
        # typical line: '   1 a      eigenvalue = -9.081551313859354D-02     i. e.  im 1549.46 cm**(-1)'
        if column[len(column)-3]=="im":
            freq=-freq  # convert imaginary 'im freq' to '-freq'

        spectra[totf].freq=freq
        column=split(infile.readline()) # should be '     atom          X             Y             Z"
        if column[0]<>'atom':
            sys.stderr.write("Parse error: expected 'atom', got '"+column[0]+"'\n")
            sys.exit(2)

        for i in range(0,len(atoms)):
            column=split(infile.readline())
            spectra[totf].x.append(float(column[2])) # x-displacement in col 2
            spectra[totf].y.append(float(column[3])) # y-displacement in col 3
            spectra[totf].z.append(float(column[4])) # z-displacement in col 4

        tempmass=split(infile.readline())
        if tempmass[0]<>'reduced':
            sys.stderr.write("Parse error: expected 'reduced', got '"+tempmass[0]+"'\n")
            sys.exit(2)
        spectra[totf].redmass=float(tempmass[3])

        spectra[totf].intensity=0.0 # not calculated

            
        totf=len(spectra)
        

    return spectra

    
#****************
# def GetAoforceSpectra()
#
# Extracts IR-spectrum data from 'infile', expects normal aoforce output format
#*********
def GetAoforceSpectra(infile):

    found=0
    line='foo'
    tempfreq=[]

    while line:
        if (len(tempfreq)>0):
            if tempfreq[0]=='frequency':
                found=1
                break   
        line=infile.readline()
        tempfreq=split(line)

    if not(found):
        sys.stderr.write("Error: No frequencies found\n")
        sys.exit(2)

    
    spectra=[]
    
    totf=0
    while 1:

        for i in range(1,len(tempfreq)):
            spectra.append(SpectraClass())

            spectra[totf+i-1].sym="-"   #   Not extracted yet!

            tempfreq[i]=string.replace(tempfreq[i],'i','-',1) # convert imaginary 'ifreq' to '-freq'
            spectra[totf+i-1].freq=float(tempfreq[i]) # frequencies in cols. 2-7

        infile.readline()
        infile.readline()
        infile.readline()
        infile.readline()
        infile.readline()
        tempint=split(infile.readline())
        if tempint[0]<>'intensity':
            sys.stderr.write("Parse error: expected 'intensity', got '"+tempint[0]+"'\n")
            sys.exit(2)

        for i in range(2,len(tempint)):
            spectra[totf+i-2].intensity=float(tempint[i]) # intensities in cols. 3-8


        infile.readline()
        infile.readline()
        infile.readline()
        infile.readline()

        tempx=[]
        tempy=[]
        tempz=[]
        
        for i in range(0,len(atoms)):
    	    tempx.append(infile.readline())
    	    tempy.append(infile.readline())
    	    tempz.append(infile.readline())

        for i in range(0,len(atoms)):
            xline=split(tempx[i])
            yline=split(tempy[i])
            zline=split(tempz[i])        
               
            for j in range(1,len(yline)): 
                spectra[totf+j-1].x.append(float(xline[j+2])) # x-displacements in cols. 4-9
                spectra[totf+j-1].y.append(float(yline[j]))   # y-displacements in cols. 2-7
                spectra[totf+j-1].z.append(float(zline[j]))   # z-displacements in cols. 2-7        
        
        infile.readline()
        temprm=split(infile.readline())
        if temprm[0]<>'reduced':
            sys.stderr.write("Parse error: expected 'reduced', got '"+temprm[0]+"'\n")
            sys.exit(2)
        
        for i in range(2,len(temprm)): # reduced masses in cols. 3-8
            spectra[totf+i-2].redmass=float(temprm[i])
                
        infile.readline()            
        infile.readline()
        infile.readline()
        infile.readline()
        totf=len(spectra)
        
        tempfreq=split(infile.readline())
        if (len(tempfreq)==0) or (tempfreq[0]<>'frequency'):
            break

    return spectra


#***************
# def RemoveZeroFreqs()
#
# This function removes all modes with frequency 0 cm**-1
# (They are not vibrations, they are the translations)
#****************
def RemoveZeroFreqs(spectra):
    i=0
    while i<len(spectra):
        if spectra[i].freq==0:
            del spectra[i]
        else:
            i=i+1

    return spectra


#*****************
# def ScaleDisps()
#
# This function scales all the displacement coordinates with the factor 'scale'
# Nice if your vibration normally is too small/large
#**************************
def ScaleDisps(scale, spectra):

    for i in range(0,len(spectra)):
        for j in range (0,len(atoms)):
            spectra[i].x[j]=spectra[i].x[j]*scale
            spectra[i].y[j]=spectra[i].y[j]*scale
	    spectra[i].z[j]=spectra[i].z[j]*scale
    
    return spectra


#***********
# def PutHCHeader()
# 
# defines the first lines of the output file.
#**************************
def PutHCHeader(outfile):    
    outfile.write("; Created by aoforce2hc "+VERSION+" ("+DATE+").\n")
    outfile.write("; All possible zero-frequencies from aoforce output ")
    if zeros:
        outfile.write("kept.\n")
    else:
        outfile.write("removed.\n")
    outfile.write("; Displacement scaling factor: %f\n\n" % (scale))
    outfile.write("; Calculation method as specified in script is just for looks.\n\n") 
    
    outfile.write("calculation-method = SemiEmpirical\n")   # HC needs something!
    outfile.write("semi-empirical-method = AM1\n")          # Why not...

        
#***************
# def PutHCSpectra()
#
# Writes IR-spectrum data file 'outfile' in G94 format
#************************'
def PutHCSpectra(outfile, spectra):

    totf=len(spectra)
    tota=len(atoms)
    
    outfile.write("ir-band-count = "+`totf`+"\n")
    for i in range(0,totf):
        outfile.write("ir-frequency(%i) = %11.6f\n" % (i+1,spectra[i].freq))
    for i in range(0,totf):
        outfile.write("ir-intensity(%i) = %11.6f\n" % (i+1,spectra[i].intensity))
    for i in range(0,totf):
        outfile.write("ir-normal-mode(%i) = " % (i+1))
        for j in range (0,tota-1):
            outfile.write('%8.6f, %8.6f, %8.6f,' % (spectra[i].x[j], spectra[i].y[j], spectra[i].z[j]))
        outfile.write('%8.6f, %8.6f, %8.6f\n' % (spectra[i].x[tota-1], spectra[i].y[tota-1], spectra[i].z[tota-1]))
    outfile.write("vibrational-mode = 1")
            
        
def outit(infile, outfile):
    infile.close()
    outfile.close()
   
   
#*****
#************
# Main
#*************
#***********

print "\naoforce2hc "+VERSION+"\n"

infile, outfile, zeros, scale=init()    

#first, decide if we have a normal aoforce output or a $les output:
les=0   # les=0; normal output
        # les=1; les output
line=infile.readline()
while line:
    if string.find(line,"LOWEST EIGENVALUE SEARCH")<>-1:
        les=1
        break
    line=infile.readline()
infile.seek(0)  # reset file

print "Extracting atom data..."    
atoms=GetAoforceAtoms(infile)

if les==0:
    print "Extracting spectral data (normal) ..."
    spectra=GetAoforceSpectra(infile)
elif les==1:
    print "Extracting spectral data ($les) ..."
    spectra=GetAoforceLesSpectra(infile)
    
if not(zeros):
    spectra=RemoveZeroFreqs(spectra)

spectra=ScaleDisps(scale, spectra)

PutHCHeader(outfile)

print "Converting spectral data..."
PutHCSpectra(outfile, spectra)

outit(infile, outfile)

print "Done. HyperChem script file created."