#!/usr/bin/ruby

# Main program for getting vasp calculation condition (progress of calculation).
#
# Written by k-fleak
#

# Usage: > ruby getCalcCondition.rb [option]
# options:
#          -a => all the status (current directory)
#          -A => all the status (recursive)
#          -e => converged energy (current directory)
#          -E => converged energy (recursive)
#          -f => all energy anyway (current directory)
#          -F => all energy anyway (recursive)

require "#{ENV["HOME"]}/myRuby/parseVaspLog.rb"
require "#{ENV["HOME"]}/myRuby/fileStat.rb"

class CalcCondition
  
  attr_reader :dirAndFinLog
  attr_reader :dirAndE0
  attr_reader :dirAndAllE0
  attr_reader :dirAndNumIonicStep    
  attr_reader :dirAndNoFinLog
  attr_reader :dirAndNoFinNSW
  attr_reader :dirAndMachine
  attr_reader :lastUpdate
    
  def initialize(dirPath)
    readDirIndex(dirPath)
    searchFinLogRelax
  end
  
  
  
  private
  
  def readDirIndex(dirPath)
    @dirIndex = Array.new()
    Dir::glob(dirPath).each do |dire|
      if FileTest::directory?(dire) then
        if FileTest.exist?("./#{dire}/INCAR") and FileTest.exist?("./#{dire}/POTCAR") and FileTest.exist?("./#{dire}/machine")  then
          @dirIndex.push(dire)
        end
      end
    end
  end
  
  def searchFinLogRelax
    @dirAndFinLog = Hash.new
    @dirAndNoFinLog = Hash.new
    @dirAndE0     = Hash.new
    @dirAndAllE0 = Hash.new
    @dirAndNumIonicStep    = Hash.new
    @dirAndNoFinNSW = Hash.new
    @dirAndMachine = Hash.new
    @lastUpdate = Hash.new
    @dirIndex.each do |di|
      dir = Dir.open(di)
      ptmp = Array.new()
      etmp = Array.new()
      etmp2 = Array.new()
      ntmp = Array.new()
      ptmp2 = Array.new()
      ntmp2 = Array.new()      
      utmp = Array.new()
      while filename = dir.read
        if /log.relax/ =~ filename then
          parse = ParseVaspLog.new("./#{di}/#{filename}")
          if parse.getConvergence == "TRUE" then
            ptmp.push(filename)
            ntmp.push(parse.getNumIonicStep)            
            etmp.push(parse.getFinalE0)
          else
            ptmp2.push(filename)
            ntmp2.push(parse.getNumIonicStep)            
          end
          etmp2.push(parse.getFinalE0)
          fstat = FileStat.new("./#{di}/#{filename}")
          utmp.push( fstat.getLastUpdate )
        elsif /log/ =~ filename then
          parse2 = ParseVaspLog.new("./#{di}/#{filename}")
          ptmp2.push("log       ")
          fstat = FileStat.new("./#{di}/#{filename}")
          utmp.push( fstat.getLastUpdate )
          if parse2.getNumIonicStep then
            ntmp2.push(parse2.getNumIonicStep)
          else
            ntmp2.push("-")
          end
        end
      end
      @dirAndFinLog[di] = ptmp[(ptmp.size) - 1]
      @dirAndNoFinLog[di] = ptmp2[(ptmp2.size) - 1]      
      @dirAndE0[di]     = etmp[(etmp.size) - 1]
      @dirAndAllE0[di]  = etmp2[(etmp2.size) - 1]
      @dirAndNumIonicStep[di] = ntmp[(ntmp.size) - 1]
      @dirAndNoFinNSW[di] = ntmp2[(ntmp2.size) - 1]
      @lastUpdate[di] = utmp[(utmp.size) - 1]
      machine = open("./#{di}/machine", "r")
      mtmp = Array.new
      machine.each do |mac|
        unless mac.strip.empty? then
          mtmp.push(mac.strip)
        end
      end
      @dirAndMachine[di] = mtmp[0]
    end
  end
    
  
  
end


if /-a/ =~ ARGV[0] or /-A/=~ ARGV[0] then
  if /-a/ =~ ARGV[0] then
    cc = CalcCondition.new("./*")
  else
    cc = CalcCondition.new("./**/")
  end
  print "\n"
  print sprintf("%12.12s %14.16s %18.12s %16.11s %18.10s %30.16ss", "Dir", "Fin?", "Log-file", "NSW", "machine", "Last update")
  print "\n", "\n"
  cc.dirAndE0.sort.each do |daf|
    dir = daf[0]
    fin = "FALSE"
    if daf[1] then
      fin = "TRUE"
      log = cc.dirAndFinLog[dir]
      nsw = cc.dirAndNumIonicStep[dir]
    else
      log = cc.dirAndNoFinLog[dir]
      nsw = cc.dirAndNoFinNSW[dir]
    end
    mac = cc.dirAndMachine[dir]
    lastUp = cc.lastUpdate[dir]
    print sprintf("%12.12s %14.16s %18.12s %16.11s %18.10s %30s", dir, fin, log, nsw, mac, lastUp)
    print "\n"
      
  end
elsif /-e/ =~ ARGV[0] or /-E/ =~ ARGV[0] then
  if /-e/ =~ ARGV[0] then
    cc = CalcCondition.new("./*")
  else
    cc = CalcCondition.new("./**/")
  end
  cc.dirAndE0.sort.each do |dae|
    lineArr = dae[0].strip.split(/\//)
    dir = lineArr[(lineArr.size) - 1]
    energy = dae[1]
    if energy then
      print dir, "     ", energy, "\n"
    end
  end
elsif /-f/ =~ ARGV[0] or /-F/ =~ ARGV[0] then
  if /-f/ =~ ARGV[0] then
    cc = CalcCondition.new("./*")
  else
    cc = CalcCondition.new("./**/")
  end
  cc.dirAndAllE0.sort.each do |dae|
    lineArr = dae[0].strip.split(/\//)
    dir = lineArr[(lineArr.size) - 1]
    energy = dae[1]
    if energy then
      print dir, "     ", energy, "\n"
    end
  end
end