#!/usr/bin/ruby

# Output position and force.
#
# First written by k-fleak
#   Time-stamp: <05/10/20>
#
# Usage:
# > ruby makePosForce.rb [option] [filename]
#
# Option:
# -c: Output HF-forces only for displaced atom.
# -m: Output atom positions and corresponding forces for Maple script "PlotForce".

home = "."

require "#{home}/parseVasprun.rb"
require "#{home}/linearAlgebra.rb"
require "#{home}/parseDISP.rb"
require "#{home}/printData.rb"
require "getopts"

class ForceForMaple
  include LinearAlgebra
  include PrintData
  
  attr_reader :force
  attr_reader :position
  
  def initialize(filename="vasprun.xml")
    @vasprun = ParseVasprun.new(filename)
  end
  
  def start
    setData()
    setAbsPosition()
    printPositionForce()
  end

  protected
  
  def setData
    @basis = @vasprun.getBasis()
    @position = @vasprun.getPosition()
    @force = @vasprun.getForce()
  end    
  
  def setAbsPosition
    @absPos = multiply(@position,@basis)
  end
  
  private
  
  def printPositionForce
    @absPos.each_with_index do |a,i|
      printArray(a, @force[i])
      printf("\n")
    end
  end 
end


class ForceRead < ForceForMaple
  def initialize(filename)
    super(filename)
    @force = @vasprun.getForce
  end

  def printForce(index)
    printArray(@force[index.to_i-1])
    printf("\n")
  end
end


class ForceForCheck
  def initialize
    @disp = ParseDISP.new()
    @atomIndex = @disp.getAtomIndex
  end

  def start
    ARGV.each_with_index do |filename,i|
      force = ForceRead.new(filename)
      force.printForce(@atomIndex[i])
    end
  end
end


unless getopts("c", "m")
  $stderr.puts("#{ARGV.shift}: illegal option\n")
  exit
end

case 
when $OPT_c
  forceCheck = ForceForCheck.new()
  forceCheck.start()
when $OPT_m
  forceMaple = ForceForMaple.new(ARGV.shift)
  forceMaple.start()
end