#! /usr/bin/env ruby

require "rubygems"
gem "crystalcell"
require "crystalcell"
require "vasputils"
require 'thor'

### Command template
class PoscarCommand < Thor
  desc "lattice", "show lattice constants"
  def lattice_constants(* args)
    puts "      a,       b,       c,  alpha,   beta,  gamma,   volume, file"
    args.each do |file|
      axes = VaspUtils::Poscar.load_file(file).axes
      latticeconstants = CrystalCell::LatticeAxes.axes_to_lc(axes)
      printf("%7.4f, %7.4f, %7.4f, %6.2f, %6.2f, %6.2f, ",
        *(latticeconstants))
      printf("%8.3f, ", CrystalCell::Cell.new(axes).calc_volume) # show volume
      puts file
    end
  end

  ## 
  desc "distance poscar_file id1 id2", "distance between atoms"
  def distance(* args)
    if args.size > 3
      puts "USAGE: poscar distance poscar_file [id1] [id2]"
      puts "Note: id count from 1"
      exit
    end

    poscar = VaspUtils::Poscar.load_file(args.shift)
    id = args.shift
    atom1s = [id.to_i]
    atom1s = 1..poscar.positions.size if id == nil

    id = args.shift
    atom2s = [id.to_i]
    atom2s = 1..poscar.positions.size if id == nil

    printf("atom1, atom2, %8s, %8s\n", "distance", "periodic_distance")
    atom1s.each do |atom1|
      atom2s.each do |atom2|
        index1 = atom1 - 1
        index2 = atom2 - 1

        cell = poscar.to_cell
        positions = cell.positions
        d =  cell.distance(positions[index1], positions[index2]) #not periodic

        pcell = cell.to_pcell
        pd = pcell.nearest_distance(positions[index1], positions[index2]) #not periodic
        printf("%5i, %5i, %8.5f, %8.5f\n", atom1, atom2 , d, pd)
      end
    end
  end

  desc "substitute POSCAR [elem1 elem2]", "Generate poscar with substituted elements."
  long_desc " When elem1 and elem2 are empty, show elements in POSCAR." +
    " When elem2 already exist in POSCAR, merge the elements."
  option :atmic_radii, desc: "Lattice constants from atomic radii."
  def substitute(* args)
    usage = "Usage: #{File.basename("#{__FILE__}")} POSCAR old_elem new_elem"

    poscar, elem1, elem2 = * args
    old_poscar = VaspUtils::Poscar.load_file(poscar)

    unless elem1 && elem2
      num = old_poscar.elements.size
      width = 7
      printf("%7s" * num, * old_poscar.elements)
      puts
      printf("%7d" * num, * old_poscar.nums_elements)
      puts
      exit
    end

    new_poscar = old_poscar.substitute(elem1, elem2)

    sum_old_radius = 0.0
    old_poscar.elements.size.times do |i|
      radius = CrystalCell::Element.atomic_radius( old_poscar.elements[i])
      sum_old_radius+= radius * old_poscar.nums_elements[i]
    end
    sum_new_radius = 0.0
    new_poscar.elements.size.times do |i|
      radius = CrystalCell::Element.atomic_radius( new_poscar.elements[i])
      sum_new_radius+= radius * new_poscar.nums_elements[i]
    end

    new_poscar.axes = old_poscar.axes.map{|axis| axis.map{|v| v * sum_new_radius/sum_old_radius}}
    new_poscar.dump($stdout)
  end

  ## POSCAR から得られる情報から、povray objects を生み出す povray 形式ファイルを生成。
  ## 描画は行わない。外部でシェルスクリプトなどで実行すべき。
  ##
  ##USAGE:
  ## poscar2pov.rb POSCAR -e 'Li,O'
  ##   POTCAR なしでも元素指定可能。POTCAR より優先。
  ##   なお、POTCAR も -e 指定もない、すなわち元素が分からない状態ではプログラムは実行されない。
  ##
  ## poscar2pov.rb POSCAR -b "Li,Ge,0.0,1.0,Ge,Ge,2.0,3.0"
  ##   以下の2つの条件の和集合となる連結の描画情報を示す POSCAR.bonds.inc を生成。
  ##   - Li-Ge間で距離 0.0〜1.0
  ##   - Ge-Ge間で距離 2.0〜3.0
  #desc "povray", "generate povray objects."
  #option :tolerance, desc: 'Tolerance for periodic shown atom.' #  '-t tolerance'
  #option :bonds, desc: 'Generate bonds matching conditions.',
  #  long_desc: 'E.g., -b "Li,Ge,0.0,1.0" indicateing two conditions'
  #  #'-b conditions',
  #def povray(* args)
  #  poscar = args[0] || 'POSCAR'
  #  tolerance = options[:tolerance].to_f || 0.0
  #
  #  cell = VaspUtils::Poscar.load_file(poscar).to_cell(CrystalCell::Povray::Cell)
  #  puts cell.atoms_to_povs(tolerance).join
  #  puts cell.lattice_to_povs.join
  #  if options[:bonds]
  #    elem0, elem1, min, max = (options[:bonds].split(','))
  #    puts cell.bonds_to_povs(elem0, elem1, min.to_f, max.to_f).join
  #  end
  #end

  desc 'png4in1 POSCAR_file', 'Generate a PNG file from 4 view points.'
  def png4in1(*args)
    poscar_name = args.shift
    cell = VaspUtils::Poscar.load_file(poscar_name).to_cell
    povray = CrystalCell::Povray.new(cell: cell)
    povray.set_axes([-1.0, -1.0, 0.0])
    povray.shoot_4in1(poscar_name)
  end

  desc 'snapgeomopt [options]', 'Generate POSCARS of each ionic steps'
  option :xdatcar, desc: "Use XDATCAR in spite of vasprun.xml"
  long_desc "Note that XDATCAR doesn't preserve cell change."
  def snapgeomopt
    if options[:xdatcar]
      xdatcar = VaspUtils::Xdatcar.load_file("./XDATCAR")
      positions_list = xdatcar.steps_positions
      elements = xdatcar.elements
      nums_elements = xdatcar.nums_elements
      bases = Array.new( positions_list.size).fill(xdatcar.axes)
    else
      xml = VaspUtils::VasprunXml.load_file("./vasprun.xml")
      bases =  xml.bases
      positions_list = xml.positions_list
      elements = xml.elements.uniq
      nums_elements = xml.nums_elements
    end

    num_iteration = positions_list.size
    width = num_iteration.to_s.size

    bases.size.times do |i|
      prefix = sprintf("%0#{width}d", i +1)
      File.open(prefix + ".vasp", "w") do |io|

        poscar = VaspUtils::Poscar.new(
          hash = {
            :comment            => "Generated by poscar_snapgeomopt, #{prefix}/#{num_iteration}",
            :scale              => 1.0,
            :axes               => bases[i],
            :elements           => elements,
            :nums_elements      => nums_elements,
            :selective_dynamics => false,
            :direct             => 'direct',
            :positions          => positions_list[i],
          }
        )
        poscar.dump(io)
      end
    end
  end

  desc "vaspdir", "Generate vaspdir"
  option :wavelength, :required => true,  desc: 'Wavelength in real space for KPOINTS.'
  def vaspdir
    system "potcar generate --poscar > POTCAR"
    system "incar generate base singlepoint spin2 metal_geomopt --enmax130 > INCAR"
    system "kpoints generate --length=#{options[:wavelength]} > KPOINTS"
  end

  desc "vaspgeomopt", "Generate vaspdir"
  option :wavelength, :required => true,  desc: 'Wavelength in real space for KPOINTS.'
  def vaspgeomopt
    poscar = VaspUtils::Poscar.load_file('POSCAR')
    es = poscar.elements
    ns = poscar.nums_elements
    dirname = ''
    poscar.elements.size.times do |i|
      dirname += es[i]
      dirname += ns[i].to_s
    end
    FileUtils.mkdir_p dirname
    subdir = dirname + '/geomopt00'
    FileUtils.mkdir_p subdir

    FileUtils.cp("POSCAR", subdir)
    system "potcar generate --poscar > #{subdir}/POTCAR"
    system "incar generate base full_relax spin2 metal_geomopt --enmax130=#{subdir}/POTCAR > #{subdir}/INCAR"
    system "kpoints generate --length=#{options[:wavelength]} > #{subdir}/KPOINTS"
  end
end

PoscarCommand.start(ARGV)