#!/usr/bin/env python

"""
molfreq.py: Print Jaguar frequencies and normal modes in a format suitable
for viewing with Molden.

Usage:
  1. molfreq.py [Jaguar output file] > freq.molf
  2. molden freq.molf

Copyright 2000, Richard P. Muller and William A. Goddard, III. All Rights
Reserved. 

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 2
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

Version History

0.1    01/01/00   Original version
0.2    03/01/00   Fixed bug related to intensities and symmetries
0.3    12/04/00   Fixed bug related to reduced masses.
0.4    08/22/02   Cleaned code slightly
0.5    06/06/03   Fixed a bug related to the force constant line
                  in Jaguar v4.2

Things to do:
* Put the frequency stuff into the MOLF format in Pistol

"""

import sys

def molfreq(filename):
    mol,freqlist = rd(filename)
    #outname = 'freq.molf'
    #outfile = open(outname,'w')

    print '[Molden Format]'
    print '[FREQ]'
    for freq in freqlist:
        print '%20.8f ' % freq
    
    print '[FR-COORD] Angs' 
    for (sym,x,y,z) in mol:
        print '%s %10.4f %10.4f %10.4f ' % (sym,x/0.52918,
                                                y/0.52918,z/0.52918)

    #It's really bone-headed to do this in two steps...
    print '[FR-NORM-COORD]'
    nfreq = len(freqlist)
    for i in range(nfreq):
        freq_num = i+1
        print 'vibration %d ' % freq_num
        mode = getmode(freq_num,filename)
        printmode(mode)
    return    

def cleansym(sym):
    # This function strips off the garbage (everything
    # after and including
    # the first non-letter) in an element name.
    import re,string
    pat = re.compile('[^a-zA-Z]')
    newsym = pat.split(sym)[0]
    return newsym

def rd(filename):
    import re,string
    geo = re.compile('[new|Input|Symmetrized|final] geometry:')
    freq = re.compile('^\s*frequencies')
    freqstart = re.compile('start of program freq')

    freqlist = []
    
    file = open(filename,'r')
    while 1:
        line = file.readline()
        if not line: break
        if geo.search(line):
            file.readline()
            file.readline()
            mol = readonegeo(file)
        if freq.search(line):
            words = string.split(line)
            for word in words[1:]:
                freqlist.append(float(word))
        if freqstart.search(line):
            freqlist = []
    file.close()
    return mol,freqlist

def readonegeo(file):
    import string
    mol = []
    while 1:
        line = file.readline()
        if not line: break
        words = string.split(line)
        if len(words) != 4: break
        sym = words[0]
        x = float(words[1])
        y = float(words[2])
        z = float(words[3])
        sym = cleansym(sym)
        mol.append((sym,x,y,z))
    return mol

def getmode(modenum,filename):
    import re,string
    freq = re.compile('^\s*frequencies')
    freqstart = re.compile('start of program freq')
    file = open(filename,'r')

    while 1:
        line = file.readline()
        if not line:
            break
        elif freqstart.search(line):
            count = 1
        elif freq.search(line):
            words = string.split(line)
            number = len(words) - 1
            if count <= modenum < count + number:
                recnum = modenum - count + 2
                mode = getonemode(recnum,file)
                file.close()
                return mode
            else:
                count = count + number
    print 'Warning: could not find mode number %d' % modenum
    file.close()
    return []

def getonemode(recnum,file):
    import string,re
    mode = []
    red_mass = re.compile('reduc. mass')
    intens = re.compile('intensities')
    symmet = re.compile('symmetries')
    forceconst = re.compile('force const')
    while 1:
        line = file.readline()
        # Skip lines we don't want
        if red_mass.search(line) or intens.search(line) \
           or symmet.search(line) or forceconst.search(line):
            continue
        if not line: break
        words = string.split(line)
        if len(words) < 2: break
        x = float(words[recnum])
        line = file.readline()
        if not line:
            print 'Warning: should not break here!'
            break
        words = string.split(line)
        if len(words) < 2: break
        y = float(words[recnum])
        line = file.readline()
        if not line:
            print 'Warning: should not break here!'
            break
        words = string.split(line)
        if len(words) < 2: break
        z = float(words[recnum])

        mode.append((x,y,z))
    return mode

def printmode(mode):
    for m in mode:
        print '%10.4f %10.4f %10.4f' % m
    return


if __name__ == '__main__':
    filename = sys.argv[1] 
    molfreq(filename)