#!/usr/bin/perl -ws

our ($h,$help,$func,$debug);
if ($h || $help) {
  (my $program = $0) =~ s/^.*[\/\\]//;
  print <<HELP;
Usage: $program file.xyz [files.xyz]
Dependences: perl, dftd3

Run Grimme dftd3 program with BJ-damping 
and add Edisp correction (kcal/mole) into 2nd string of xyz-files

Options:
-finc=functional  Functional name in TurboMole style (-finc=pbe by default)
-debug
HELP
  if ($help) {
    print <<HELP;
    
DFT-D3: 
http://www.thch.uni-bonn.de/tc/index.php?section=downloads&subsection=DFT-D3
Please cite
S. Grimme, J. Antony, S. Ehrlich and H. Krieg,  J. Chem. Phys. 132 (2010), 154104
S. Grimme, S. Ehrlich and L. Goerigk, J. Comput. Chem. 32 (2011), 1456-1465
HELP
  }
  exit;
}

#DFT-D3 program (http://www.thch.uni-bonn.de/tc/downloads/DFT-D3/data/dftd3.tgz)
my $dftd3_x = '/usr/local/bin/dftd3';
if (! -x $dftd3_x) {
  die "No $dftd3_x!\n";
}

$func ||= 'pbe';

# Known functionals (names in TurboMole style)
# http://www.thch.uni-bonn.de/tc/downloads/DFT-D3/functionalsbj.html
my @func = qw (
b-p b-lyp revpbe rpbe b97-d pbe rpw86-pbe b3-lyp tpss hf tpss0 pbe0 hse06 
revpbe38 pw6b95 b2-plyp dsd-blyp dsd-blyp-fc bop mpwlyp o-lyp pbesol bpbe opbe 
ssb revssb otpss b3pw91 bh-lyp revpbe0 tpssh mpw1b95 pwb6k b1b95 bmk cam-b3lyp 
lc-wpbe b2gp-plyp ptpss pwpb95 hf/mixed hf/sv hf/minis b3-lyp/6-31gd hcth120 
dftb3 pw1pw pwgga hsesol hf3c hf3cv);

# Some synonims of functionals
$func = lc($func);
my %synonims = qw(
b3lyp  b3-lyp
blyp   b-lyp
plyp   o-lyp
pbe1   pbe0
);
if (exists $synonims{$func}) {
  $func = $synonims{$func};
}

if (! grep {$func eq $_} @func) {
  die "Known functionals: @func\n";
}

foreach my $file (@ARGV) {
  my @mols = read_molden($file);
  
  # Remove old Edisp
  foreach my $mol (@mols) {
    while (1) {
      $mol->[0] =~ s/Edisp\s+-?\d+(\.\d+)?\s*// or last;
    }
    $mol->[0] =~ s/\s+$//;
  }
  
  if (@mols > 1) {
    my $tmp = 'dftd3.xyz';
    foreach my $mol (@mols) {
      write_molden($mol, $tmp);
      my $D3 = D3correction($tmp);
      $mol->[0] .= "  Edisp $D3";
    }
    unlink $tmp;
    write_molden(@mols,$file);
  }
  else {
    my $D3 = D3correction($file);
    $mols[0][0] .= "  Edisp $D3";
    write_molden($mols[0], $file);
  }
}

sub D3correction {
  my $file = shift; # filename of xyz as input
  my $D3;           # return value
  my @args = ($dftd3_x, $file, '-func', $func, '-bj');
  print "Run `@args`\n" if $debug;
  open DFTD3, '-|', @args or do {warn "Can't run @args : $!\n"; return undef};
  while (<DFTD3>) {
    print if $debug;
    if (/^\s*Edisp \/kcal,au:\s+(\S+)/) {
      $D3 = $1;
    }
  }
  close DFTD3;
  if (! defined($D3)) {
    die "Problems in subroutine D3correction\n";
  }
  return sprintf("%.2f", $D3);
}

sub read_molden {
 ############################################################################
 ## Молекула -- ссылка на массив, 0-й элемент -- свойства, следующие - атомы.
 ## Свойства -- ссылка на хэш с ключами Energy, Symmetry
 ## Атом -- ссылка на массив [atom, x, y, z, ppm] (ppm может не быть).
 ##
 ## Читает xyz. Параметры - имена xyz-файлов. Если параметров нет, то <>.
 ## Возвращает массив найденных молекул.
 ############################################################################
  local @ARGV = @_ ? @_ : @ARGV;
  my $num = qr/-?\d+(?:\.\d+)?/;
  my @mols;
  my $line;
  LOOP:
  while ($line || defined($line = <>)) {
    #print $line;
    if ($line =~/^\s*(\d+)\s*$/) {
      my @mol;
      my $N = $1;
      last LOOP if eof();
      next LOOP if eof(ARGV);
      $line = <>;
      chomp $line;
      $mol[0] = $line;
      for (my $i=1; $i<=$N; $i++) {
        last LOOP if eof();
        next LOOP if eof(ARGV);
        $line = <>;
        #print $line;
        if ($line =~ /^\s*([A-Z]{1,2})\s+($num)\s+($num)\s+($num)\s*(.*)/io) {
          $mol[$i] = [$1,$2,$3,$4,$5];
        } else {
          next LOOP;
        }
      }
      push @mols, \@mol;
      last LOOP if eof();
    } else {
      undef $line;
    }
  }
  return @mols;
}

sub write_molden {
  my $file = pop @_;
  open F, '>', $file or do {warn "Can't write to $file: $!\n"; return};
  foreach my $mol (@_) {;
    my $N = $#{$mol};
    print F " $N\n$mol->[0]\n";
    for (my $i=1; $i<=$N; $i++) {
      printf F " %-2s %12.8f %12.8f %12.8f", $mol->[$i][0],$mol->[$i][1],$mol->[$i][2],$mol->[$i][3];
      print F "  $mol->[$i][4]" if defined $mol->[$i][4];
      print F "\n";
    }
  }
  close F;
}