#!/usr/bin/ruby

# Main program for preparing VASP calculation in CMC ansatz.
#
# Written by k-fleak
#

require 'ftools'

require "#{ENV["HOME"]}/myRuby/atom.rb"
require "#{ENV["HOME"]}/myRuby/cell.rb"
require "#{ENV["HOME"]}/myRuby/parseCifP1.rb"


class PrepareVASPCalc
  
  def initialize
    searchPOSCAR
    makeDir
  end
  
  private
  
  def makeDir
    @dirName.each do |dir|
      unless FileTest.exist?(dir) then
        Dir::mkdir(dir, 0755)
        element = getElementPOSCAR("./POSCAR.#{dir}")
        potcar = getPOTCARname(element)
        File.copy "./POSCAR.#{dir}", "./#{dir}/POSCAR"
        File.copy "./INPUTS/#{potcar}", "./#{dir}/POTCAR"
        File.copy "./INPUTS/INCAR", "./#{dir}/INCAR"
        File.copy "./INPUTS/KPOINTS", "./#{dir}/KPOINTS"
        unless FileTest.exist?("./POTCAR.cmc") then
          Dir::mkdir("./POTCAR.cmc", 0755)
        end
        unless FileTest.exist?("./POSCAR.cmc") then
          Dir::mkdir("./POSCAR.cmc", 0755)
        end        
        File.copy "./INPUTS/#{potcar}", "./POTCAR.cmc/POTCAR.#{dir}"
        delComment("./POSCAR.#{dir}", "POSCAR.tmp")       
        File.move "POSCAR.tmp" , "./POSCAR.cmc/POSCAR.#{dir}"
        File.delete("./POSCAR.#{dir}")        
      else
        File.delete("./POSCAR.#{dir}")
      end
    end
  end
  
  def getPOTCARname(element)
    tailName = "POTCAR."
    if element[0] then
      element.each do |elem|
        tailName = tailName + elem
      end
    else
      tailName = "POTCAR"
    end
    return tailName
  end
  
  def delComment(filename, outname)
    input = open(filename)
    output = Array.new
    input.each do |inp|
      if /#/ =~ inp then
        ptmp = inp.split(/#/)
        pos = ptmp[0] 
        output.push(["#{pos}  \n"])
      else
        output.push(["#{inp}"])
      end
    end
    outfile = open(outname, "w")
    output.each do |outs|
      outfile.write(outs)
    end
    outfile.close
  end
  
  def getElementPOSCAR(filename)
    input = open(filename, "r")
    flag = 0
    elemArray = Array.new()
    input.each do |inp|
      if flag == 1 then
        if /#/ =~ inp then
          lineArr = inp.strip.split(/#/)
          etmp = lineArr[lineArr.size - 1]
          elemArray.push(etmp.strip.split(/-/)[0])
        else
          elemArray.push(nil)
        end
      end
      if /Direct/ =~ inp then
        flag = 1
      end
    end
    elemArray.uniq!
    return elemArray    
  end
  
  def searchPOSCAR
    @posFileArray = Array.new()
    @dirName = Array.new()
    Dir.entries("./INPUTS/").each do |dir|
      if FileTest.file?("./INPUTS/#{dir}") and /POSCAR/ =~ dir then
        if /.cif/ =~ dir then
          newName = dir.gsub(/.cif/, '')
          cif2poscar(dir, newName)
          @posFileArray.push(newName)
          @dirName.push(newName.gsub(/POSCAR./, ''))
        elsif
          @posFileArray.push(dir)
          @dirName.push(dir.gsub(/POSCAR./, ''))
          File.copy "./INPUTS/#{dir}", "./#{dir}"
        end
      end
    end
  end
  
  def cif2poscar(filenameOrig, filenameNew)
    cif = ParseCifP1.new("./INPUTS/#{filenameOrig}")
    position = cif.getPosition
    axis = cif.getAxis
    comment = cif.getComment
    name = cif.getAtomName
    atomArray = Array.new
    position.each_with_index do |x, i|
      atom = Atom.new(x, axis, name[i])
      atomArray << atom
    end
    atoms = Cell.new(atomArray)
    poscar = atoms.getPoscar
    ftmp = open("#{filenameNew}", "w")
    poscar.each do |pos|
      ftmp.write(pos)
    end
    ftmp.close
  end
  
end

ppp = PrepareVASPCalc.new