# Parser of "vasprun.xml".
#
# First written by k-fleak
#   Time-stamp: <05/10/08>

require "rexml/document"

class ParseVasprun
  
  def initialize(filename="vasprun.xml")
    file = File.new(filename)
    @vasprun = REXML::Document.new file
  end
  
  def getForce
    @force = parseForce
    @force
  end
  
  def getNumAtom
    @numAtom = parseNumAtom
    @numAtom
  end
  
  def getBasis(state="finalpos")
    @basis = parseBasis(state)
    @basis
  end
  
  def getPosition(state="finalpos")
    @position = parsePosition(state)
    @position
  end
  
  
  private
  
  def parseNumAtom
    numAtom = @vasprun.elements["modeling/atominfo/atom"]
    numAtom
  end
      
  def parseForce
    force = Array.new
    @vasprun.elements.each("modeling/calculation/varray") do |e|
      case e.attributes["name"]
      when "forces"
        e.elements.each("v") do |v|
          tmpForce = v.text.strip.split(/\s+/)
          tmpForce.size.times do |i|
            tmpForce[i] = tmpForce[i].to_f
          end
          force.push(tmpForce)
        end
      end
    end
    force
  end
  
  def parseBasis(state)
    basis = Array.new
    @vasprun.elements.each("modeling/structure") do |e|
      case e.attributes["name"]
      when state
        e.elements.each("crystal/varray") do |v|
          case v.attributes["name"]
          when "basis"
            v.elements.each("v") do |b|
              tmpBas = b.text.strip.split(/\s+/)
              tmpBas.size.times do |i|
                tmpBas[i] = tmpBas[i].to_f
              end
              basis.push(tmpBas)
            end
          end
        end
      end
    end
    basis
  end
  
  def parsePosition(state)
    position = Array.new
    @vasprun.elements.each("modeling/structure") do |e|
      case e.attributes["name"]
      when state
        e.elements.each("varray") do |v|
          case v.attributes["name"]
          when "positions"
            v.elements.each("v") do |p|
              tmpPos = p.text.strip.split(/\s+/)
              tmpPos.size.times do |i|
                tmpPos[i] = tmpPos[i].to_f
              end
              position.push(tmpPos)
            end
          end
        end
      end
    end
    position
  end
  
end