#!/usr/bin/perl -ws

our($h,$help, $auto, $i, $B, $T, $kcal, $kJ);
(my $program = $0) =~ s/^.*[\/\\]//;

if ($h || $help) {
  print <<HELP;

Usage: $program 'averagings' file1.xyz file2.xyz ... > multifile.xyz
       cat file1.xyz file2.xyz ... | $program 'averagings'  > multifile.xyz
       $program -i file1.xyz file2.xyz ...
       $program -B file1.xyz file2.xyz ... > file.xyz
Зависимости: perl, symmetry для опции -auto (опционально)

Усреднение свойств в молденовских файлах (например, хим. сдвигов в *.xyzppm)

averagings - перечисление групп  (через пробел) 
                          атомов (через запятую и/или дефис)
Если перед номером стоит буква A, то это водороды, связанные с этим атомом
Примеры:
'3-5,7-9 1,6,10' - усреднить хим. сдвиги атомов 3,4,5,7,8,9 и (отдельно) 1,6,10
A15-A17 - усреднить хим. сдвиги протонов, связанных с атомами 15,16,17

Если вместо 'averagings' будет -auto, то автоматически будут усреднены
хим. сдвиги протонов метилов и эквивалентных атомов (для последнего действия
требуется программа symmetry).

С опцией -i, если программе даны файлы в качестве аргументов, они будут 
отредактированы по месту.

С опцией -B будет усреднение хим. сдвигов по Больцману. Во 2-й строке xyz-файлов
  д.б. энергия. Нумерация атомов во всех xyz-файлах д.б. согласованной. 
  Напечатается геометрия самого стабильного конформера с хим. сдвигами, 
  представляющими собой взвешенное среднее по всем. Веса - мольные доли,
  вычисляются из энергий.
Опции, относящиеся к -B:
-T=298    Температура
       По умолчанию энергии в a.u. (hartree)
-kkal  Энергии в ккал/моль
-kJ    Энергии в кДж/моль

HELP
  exit;
}

if ($B) {
  $T ||= 298.15;
  my $factor = 627.5; # default energy unit is a.u. (hartree)
  if    ($kcal) {$factor = 1;}
  elsif ($kJ)   {$factor = 1/4.184}
  my @mols = @ARGV ? readmolden(@ARGV) : readmolden();
  my $mol = bolzman(\@mols, $T, $factor);
  writemolden($mol);
  exit;
}

if (! $auto && ! @ARGV ) {
  die "Нужно указать, что усреднять, либо задать опцию -auto. $program -h for help\n";
}

$replace = 1 if $i;

if ($auto) {
  
  $symmetry = '/usr/local/bin/symmetry';
  $symmetry = '' unless (-x $symmetry);
  #die "Program symmetry not found\n" unless (-x $symmetry);

  # Вгоняем таблицу Менделеева
  @ATOM = qw (XX
              H                                                  He
              Li Be                               B  C  N  O  F  Ne
              Na Mg                               Al Si P  S  Cl Ar
              K  Ca Sc Ti V  Cr Mn Fe Co Ni Cu Zn Ga Ge As Se Br Kr
              Rb Sr Y  Zr Nb Mo Tc Ru Rh Pd Ag Cd In Sn Sb Te I  Xe
              Cs Ba La
                      Ce Pr Nd Pm Sm Eu Gd Tb Dy Ho Er Tm Yb Lu
                       Hf Ta W  Re Os Ir Pt Au Hg Tl Pb Bi Po At Rn
              Fr Ra Ac
                      Th Pa U  Np Pu Am Cm Bk Cf Es Fm Md No Lr
             );
  for ($i=0; $i<@ATOM; $i++) { $ATOM{$ATOM[$i]} = $i };

  $replace = 0 unless @ARGV;
  
  while (&read_molden) {

    # Усредняем протоны метильных групп
    for ($i=1; $i<=$N; $i++) {
      if ( $atom[$i] eq 'C' ) {
        $CH[$i] = [];
        for ($j=1; $j<=$N; $j++) {
          if ( $atom[$j] eq 'H' &&
               ($x[$i]-$x[$j])**2+($y[$i]-$y[$j])**2+($z[$i]-$z[$j])**2 < 1.44
              ) {
            push @{$CH[$i]}, $j;
          }
        }
      }
    }
    
    for ($i=1; $i<=$N; $i++) {
      if ($atom[$i] eq 'C' && @{$CH[$i]}==3) {
        my $aver = ($ppm[$CH[$i][0]]+$ppm[$CH[$i][1]]+$ppm[$CH[$i][2]])/3;
        $ppm[$CH[$i][0]] = $ppm[$CH[$i][1]] = $ppm[$CH[$i][2]] = $aver;
      }
    }

    if ($symmetry) {
      open TINP, ">symm.TINP" || warn "Can't write temporary file: $!\n" , next;
      print TINP " $N\n";
      for ($i=1; $i<=$N; $i++) {
        print TINP " $ATOM{$atom[$i]}   $x[$i]  $y[$i]  $z[$i]\n";
      }
      close TINP;

      $generator = undef;
      my $prim = 0.05;
      my $final = 0.2;
      open OUT, '-|', "$symmetry -primary $prim -final $final symm.TINP" or
                      warn "Can't run $symmetry: $!\n", return undef;
      #open OUT, "$symmetry -final 0.009999 symm.TINP |";
      while (<OUT>) {
        if (/^Molecule has the following abelian symmetry generators: (\S+)/) {
          $generator = $1;
          last;
        }
      }
      while (<OUT>) {
        if (m/^Abelian symmetrized coordinates of all atoms$/) {
          <OUT>;
          $N = <OUT>; chomp $N;
          for ($i=1; $i<=$N; $i++) {
            ($natom,$x[$i],$y[$i],$z[$i]) = split ' ', <OUT>;
            $atom[$i] = $ATOM[$natom];
          }
        }
      }
      close OUT;
      unlink 'symm.TINP';

      @symm_atoms = symm_atoms($generator);
      SYMM_ATOMS:
      foreach (@symm_atoms) {
        #print "@{$_}\n";
        my $sum = 0;
        foreach (@{$_}) {
          next SYMM_ATOMS if $ppm[$_] eq '';
          $sum += $ppm[$_];
        }
        my $aver = $sum/@{$_};
        $ppm[shift @{$_}] = $aver;
        foreach (@{$_}) {
          $ppm[$_] = '';
        }
      }
    }
    
    &write_molden;
  }
  exit;
  
  sub symm_atoms {
    my $generator = shift;
    my @symm_atoms;
    unless ($generator) {
      @symm_atoms = map {[$_]} 1..$N;
      return @symm_atoms;
    }
    my %condition = (
      'z'        => '$x[$i]==$x[$j] && $y[$i]==$y[$j] && $z[$i]==-$z[$j]',
      'x,y'      => 'abs($x[$i])==abs($x[$j]) && 
                     abs($y[$i])==abs($y[$j]) && 
                     $z[$i]==$z[$j]',
      'xy'       => '$x[$i]==-$x[$j] && $y[$i]==-$y[$j] && $z[$i]==$z[$j]',
      'x,y,z'    => 'abs($x[$i])==abs($x[$j]) && 
                     abs($y[$i])==abs($y[$j]) && 
                     abs($z[$i])==abs($z[$j])',
      'xyz'      => '$x[$i]==-$x[$j] && $y[$i]==-$y[$j] && $z[$i]==-$z[$j]',
      'xy,z'     => '($x[$i]==-$x[$j] && $y[$i]==-$y[$j] && $z[$i]==$z[$j]) || 
                     ($x[$i]==$x[$j] && $y[$i]==$y[$j] && $z[$i]==-$z[$j])',
      'xy,xz,yz' => '($x[$i]==-$x[$j] && $y[$i]==-$y[$j] && $z[$i]==$z[$j]) ||
                     ($x[$i]==-$x[$j] && $y[$i]==$y[$j] && $z[$i]==-$z[$j]) ||
                     ($x[$i]==$x[$j] && $y[$i]==-$y[$j] && $z[$i]==-$z[$j])'
    );
    my @at = @atom;
    for ($i=1; $i<=$N; $i++) {
      next unless $at[$i];
      push @symm_atoms, [($i)];
      #$at[$i] = undef;
      for ($j=$i+1; $j<=$N; $j++) {
        next unless $at[$j];
        if (exists $condition{$generator} && eval "$condition{$generator}") {
            $at[$j] = undef;
            push @{$symm_atoms[-1]}, $j;
        }
      }
    }
    return @symm_atoms;
  }
  
}


$i = 0;
foreach my $av (split ' ', shift) {
  unless ($av =~ /^(A?\d+(-A?\d+)?,)*A?\d+(-A?\d+)?$/g) {
    die "Not valid numbers: $av\n";
  }

  foreach (split /,/, $av) {
    if (/^(\d+)-(\d+)$/) { # Раскрываем интервалы (i-n)
      my @a = $1<=$2 ? $1..$2 : reverse $2..$1;
      push @{$averagings[$i]}, @a;
    }
    elsif (/^A(\d+)-A(\d+)$/) { # Раскрываем интервалы (Ai-An)
      my @a = $1<=$2 ? $1..$2 : reverse $2..$1;
      @a = map {"A$_"} @a;
      push @{$averagings[$i]}, @a;
    }
    elsif (/^A(\d+)-(\d+)$/ || /^(\d+)-A(\d+)$/) { # Проверка
      die "Not valid numbers: $_\n";
    }
    else {push @{$averagings[$i]}, $_}
  }
  $i++; 
}

#for (@averagings) {print "@{$_}\n"}
#exit;

$replace = 0 unless @ARGV;

while (&read_molden) {
  foreach (@averagings) {
    my @a = @{$_};
    if ($a[0] =~ /^A/ ) {
      @a = map {s/^A//; $_} @a;
      #print "    a @a\n";
      undef @new_a;
      foreach $A (@a) {
        for ($i=1; $i<=$N; $i++) {
          if ($atom[$i] eq 'H' && $i != $A) {
            $r = ($x[$i]-$x[$A])**2 + ($y[$i]-$y[$A])**2 + ($z[$i]-$z[$A])**2;
            push @new_a, $i if $r < 1.44;
          }
        }
      }
      @a = @new_a;
      #print "new_a @a\n";
    }
    my $sum = 0;
    foreach (@a) {
      die "Not valid number @_\n" if $_>$N;
      $sum += $ppm[$_];
    }
    $sum /= $#a+1;
    $ppm[$_] = $sum foreach @a;
  }

  &write_molden;
}


sub read_molden {
  return undef if eof();
  # 1-st string
  if (<>=~/^\s*(\d+)\s*$/) {
    ($N,$energy,@atom,@x,@y,@z,@ppm) = undef;
    $N = $1;
  } else {
    return undef;
  }
  # 2-nd string
  ($energy) = <>=~/(-?\d+\.\d+)/;
  # Geometry
  for ($i=1;$i<=$N;$i++) {
    ($atom[$i],$x[$i],$y[$i],$z[$i],$ppm[$i],undef) = split ' ', <>;
    return undef unless $atom[$i]=~/^\w\w?/ &&
                        $x[$i]=~/^-?\d+\.\d+/ &&
                        $y[$i]=~/^-?\d+\.\d+/ &&
                        $z[$i]=~/^-?\d+\.\d+/;
  }
  return $N;
}

sub write_molden {
  if (@atom && @x && @y && @z) {
    if ($replace) {
      open F, '>', "$ARGV.TMP" or die "Can't open $ARGV.TMP: $!\n";
      select F;
    }
    print " $#atom\n";
    print " Energy $energy" if $energy;
    print "\n";
    for ($i=1; $i<=$#atom; $i++) {
      if ($atom[$i] && $x[$i] && $y[$i] && $z[$i]) {
        printf " %-2s %12.8f %12.8f %12.8f",$atom[$i],$x[$i],$y[$i],$z[$i];
        printf $atom[$i] eq 'H' ? " %10.3f" : " %9.2f"  , $ppm[$i] if $ppm[$i];
        print "\n";
      }
    }
    if ($replace) {
      close F;
      rename "$ARGV.TMP", $ARGV or die "Can't rename $ARGV.TMP: $!\n";
    }
  }
}

sub bolzman {
  my @mols = @{$_[0]};
  my $T = $_[1] || 298.15;
  my $factor = $_[2] || 627.5;
  my $R = 1.987;
  @mols = grep {defined($_->[0]{Energy})} @mols;
  my $mol = $mols[0];
  foreach (@mols) {
    $mol = $_ if $_->[0]{Energy} < $mol->[0]{Energy};
  }
  $mol = copy_mol($mol);
  for (my $i=1; $i<@$mol; $i++) {
    undef $mol->[$i][4];
  }
  my @gamma = map {exp(-($_->[0]{Energy}-$mol->[0]{Energy})*$factor*1000/$R/$T)} @mols;
  my $sum = 0;
  $sum += $_ foreach @gamma;
  @gamma = map {$_/$sum} @gamma;
  #print "@gamma\n";
  ATOM:
  for (my $i=1; $i<@$mol; $i++) { # by atoms
    my $ppm = 0;
    for (my $j=0; $j<@mols; $j++) { # by mols
      next ATOM unless defined($mols[$j][$i][4]) && $mols[$j][$i][4]=~/^-?\d+(?:\.\d+)?$/;
      $ppm += $gamma[$j]*$mols[$j][$i][4];
    }
    $mol->[$i][4] = $ppm;
  }
  return $mol;
}

sub readmolden {
 ############################################################################
 ## Молекула -- ссылка на массив, 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 = <>;
      ($mol[0]{Energy}) = $line =~ /($num)/o;
      ($mol[0]{Symmetry}) = $line =~ /symm\S*\s+(\S+)/i;
      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];
          if ($mol[$i][4]=~/^($num)\s*(.*)/io) {
            $mol[$i][4] = $1;
            $mol[$i][5] = $2 if $2;
          }
        } else {
          next LOOP;
        }
      }
      push @mols, \@mol;
      last LOOP if eof();
    } else {
      undef $line;
    }
  }
  return @mols;
}

sub writemolden {
  foreach my $mol (@_) {;
    my $N = $#{$mol};
    print " $N\n";
    print " Energy $mol->[0]{Energy}" if $mol->[0]{Energy};
    print "  Point $mol->[0]{Point}" if $mol->[0]{Point};
    print "  Symmetry $mol->[0]{Symmetry}" if $mol->[0]{Symmetry};
    print "  Ellips $mol->[0]{Ellips}" if $mol->[0]{Ellips};
    print "\n";
    for (my $i=1; $i<=$N; $i++) {
      printf " %-2s %12.8f %12.8f %12.8f", $mol->[$i][0],$mol->[$i][1],$mol->[$i][2],$mol->[$i][3];
      printf " %8.2f", $mol->[$i][4] if defined($mol->[$i][4]) && $mol->[$i][4]=~/-?\d+(?:\.\d+)?/;
      print "  $mol->[$i][5]" if defined($mol->[$i][5]);
      print "\n";
    }
  }
}

sub copy_mol {
  my $mol = shift;
  my $new_mol;
  my $N = $#{$mol};
  $new_mol->[0] = $mol->[0] if $mol->[0];
  for (my $i=1; $i<=$N; $i++) {
    $new_mol->[$i] = [@{$mol->[$i]}];
  }
  return $new_mol;
}