#!/usr/bin/env python

# gennrvs.py -- script to generate NRVS spectrum from frequency calculation
# 
# Depends on cclib, Numeric/numpy, and numarray
#
# Copyright (c) Adam Tenderholt, Stanford University, 2008
#                               a-tenderholt@stanford.edu
#
# 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 WARRENTY; without even the implied warrenty of 
# MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. See the 
# GNU General Public License (http://www.gnu.org/licenses) for more 
# details.

__revision__ = "$Rev: 28 $"

from cclib.parser import ccopen
from cclib.progress import TextProgress
try:
 import Numeric as numpy
except ImportError:
 import numpy

import numarray.convolve._lineshape as lineshape
import sys
import logging
from math import sqrt, exp
from getopt import gnu_getopt

h = 6.626069e-34 #(J s)
kb = 1.38065e-23 #(J/K)
c = 2.99792458e10 #(cm^-1/s)

def usage():
  print "Usage: gennrvs filename [-f factor] [-O index] [-T temp] [-n num] [-w width]"
  print "       -f factor    frequency scaling factor (default 1.0)"
  print "       -O index     index of oxygen atom"
  print "       -T temp      temperature of the experiment"
  print "       -n num       number of frequencies (default 60)"
  print "       -w width     width of gaussians (default 16)"
  
flags, file = gnu_getopt(sys.argv[1:],"f:O:T:n:w:h",["help"])
use_temp = False

# see which, if any, command line flags have been passed
# and act accordingly

for flag in flags:
  if flag[0] == "--help" or flag[0] == "-h":
    usage()
    sys.exit(0)
  elif flag[0] == "-f":
    factor = float(flag[1])
  elif flag[0] == "-O":
    oxygen = int(flag[1])
  elif flag[0] == "-T":
    temp = float(flag[1])
    use_temp = True
  elif flag[0] == "-n":
    num = int(flag[1])
  elif flag[0] == "-w":
    width = int(flag[1])
    
if len(file) != 1:
  usage()
  sys.exit(1)
  
# if certain arguments have not been set
# via the command-line, do that now

if globals().keys().count("factor") == 0:
  factor = 1.0
if globals().keys().count("oxygen") == 0:
  oxygen = -1
if globals().keys().count("num") == 0:
  num = 60
if globals().keys().count("width") == 0:  
  width = 16

# use cclib to parse frequencies and 
# displacements from logfile

progress = TextProgress()
parser = ccopen(file[0], progress, logging.WARNING)
data = parser.parse()

if data:
  object = data
else:
  object = parser

if len(object.vibfreqs) < num:
  num = len(object.vibfreqs)

# re-open file and pull out information about 
# atomic weights used in calculation

weights = []
file = open(file[0])
for line in file:
  
  if line.find("Thermochemistry") > 0:
    lines = file.next()
    temps = file.next()
    
    line = file.next()
    while line.find("Molecular") < 0:
      
      info = line.split()
      weights.append(float(info[-1]))
      line = file.next()

#finish getting info from file

areas = []
distarray = []
indices = []

# for each frequency, calculation the intensity
# I = \frac{\sum_{Fe} m_a/N * d_a^2}{\sum_{rest} m_a * d_a^2}
# where m is the atomic mass, N is the number of iron atoms, 
# and d is the displacement of the atom in the given frequency

for i in range(len(object.vibfreqs)):
  disp = object.vibdisps[i]
  dists = reduce(numpy.add, numpy.transpose(disp**2))
  distarray.append(dists)

  #keep track of things for each frequency
  numerator = 0
  denomenator = 0
  count = 0
  
  #iterate over all atoms
  for j in range(len(object.atomnos)):
    if object.atomnos[j] == 26:
      numerator += weights[j] * dists[j]
      count += 1
      #keep track of where the irons are for displacements
      if i == 0:
        indices.append(j)
    else:
      denomenator += weights[j] * dists[j]

  area = numerator / (count * denomenator)
  if use_temp:
      exponential = (h*c*object.vibfreqs[i])/(kb*temp)
      nalpha = 1/(exp(exponential) - 1)
      areas.append(area*(nalpha+1))
  else:
      areas.append(area)
  
total_area = 0
for area in areas:
  total_area += area
print "area:", total_area

xarray = numpy.arange(0,object.vibfreqs[num - 1]*factor,1)

# for each frequency, convolute intensity 
# with a gaussian

lines = []
for i in range(num):
  lines.append(areas[i]*lineshape.gauss(xarray,width,object.vibfreqs[i]*factor))
  
# create the spectrum by "adding" together 
# gaussians created above

templines = numpy.array(lines)
yarray = reduce(numpy.add,templines)

# write spectrum to a file

file = open("spectrum.dat","w")
for i in range(len(xarray)):
  file.write("%6.3f,%9.6f\n"%(xarray[i],yarray[i]))
file.close()

# write a table of the iron displacements for
# each frequency in the logfile

file = open("displacements.txt","w")
file.write("%8s %8s " % ("  Freq.  ", " Freq.' "))
for i in range(count - 1):
  file.write("%5s " % (" Fe  "))
file.write("%5s %5s\n" % (" Fe  ", " O   "))
file.write("-" * (count*6 + 23)+"\n")

for i in range(num):
  file.write("%8.3f %8.3f "%(object.vibfreqs[i], object.vibfreqs[i]*factor))
  for index in indices[:-1]:
    file.write("%5.2f " % (sqrt(distarray[i][index])))
  if oxygen > 0:
    file.write("%5.2f %5.2f\n" % (sqrt(distarray[i][indices[-1]]), 
                                sqrt(distarray[i][oxygen])))

  else:
    file.write("%5.2f\n" % (sqrt(distarray[i][indices[-1]])))

file.close()

# finish writing table of Fe displacements

#EOF