#! @PERL@ -w
#
# Evaluates total energy and forces for a given bond dissection
#
# We assume to have the following, an $Inputdir with the following contents:
# - For each Atomic Fragment
#   - A directory with results, especially pcp.energy.all and pcp.forces.all
#   - Files called AtomicFragmentTE-Factors.dat and AtomicFragmentForces-Factors.dat
# - For Each Bond Fragment
#   - A directory with results, especially pcp.energy.all and pcp.forces.all
#   - Files called BondFragmentTE-Factors.dat and BondFragmentForces-Factors.dat

use strict;

my ($Help, $Prefix, $Inputdir);
my ($i, $j, $k, $l);
use Getopt::Long;
use File::Copy;

GetOptions("help"    => \$Help,
	"prefix=s",		=>	\$Prefix, 
	"inputdir=s"	=>	\$Inputdir);

if ($Help) { usage(); }

sub usage {
    print STDERR <<"EOF";
Usage: $0 [OPTIONS]
Evaluates the approximate expression of a many-body bond order dissection and writes pcp.energy and pcp.forces files in the inputdir
  --prefix <prefix> prefix in front of .forces.all and .energy.all
  --inputdir <inputdir> read data from this directory
  --help         		this help
EOF
    exit 1;
}

if ((!defined $Inputdir) || (!defined $Prefix)) {
	print "Inputdir and/or Prefix not specified! See --help.\n";
	exit 1;
}

# Total energy
my @TotalEnergy;
my @TotalForce;
my $energy;
my $forces;
my @tmp;

# Parse in total energy factors (1d array: Fragment)
open(TEFactors, "<$Inputdir/BondFragmentTE-Factors.dat") or die "could not open $Inputdir/BondFragmentTE-Factors.dat: $!";
my $line = <TEFactors>;
chomp($line);
close(TEFactors);
my @Factors = split(/[ \t]+/, $line);
my $lasttime = 0;

# Print parsed total energy factors
print "TEFactors:\n";
foreach $energy (@Factors) {
	print $energy."\t";
}
print "\n";

# Parse in Force factors (3d array: AtomPerFragment,AtomPerWholeMolecule,Fragment )
my @ForceFactors;
open(FORCEFactors, "<$Inputdir/BondFragmentForces-Factors.dat") or die "could not open $Inputdir/BondFragmentForces-Factors.dat: $!";
do {
	my @FragmentFactors;
	do {
		$line = <FORCEFactors>;
		chomp $line;
		my @tmp = split(/[ \t]+/, $line);
		if ( $#tmp ) {
			push @FragmentFactors, [ @tmp ];	# push a copy of the parsed line
		}
		#foreach $forces ( @tmp ) {
		#	print "\t [ $forces ]";
		#}
		$energy = $#tmp;
		#print "\t: ".$#tmp." atoms.\n";
	} until ( $energy == -1 );	# next fragment
	#print "Next Fragment!\n";
	# read over next line
	push @ForceFactors, [ @FragmentFactors ];	# push a copy of the parsed Fragment
	$line = <FORCEFactors>;
} while(<FORCEFactors>);
close(FORCEFactors);

# Print parsed Force Factors
print "\nForce Factors:\n";
for $i ( 0 .. $#ForceFactors ) {
	print "Fragment $i\n";
	for $j ( 0 .. $#{$ForceFactors[$i]} ) {
		for $k ( 0 .. $#{${$ForceFactors[$i]}[$j]} ) {
			print $ForceFactors[$i][$j][$k]."\t";
		}
		print "\n";
	}
	print "\n";
}

# check on the number of digits
my $format = "-1";
for($i=1;$i<10;$i++)
{
	my $test = "%0".$i."s";
	my $filename = sprintf($test,$i);
	if ( -e "$Inputdir/BondFragment$filename" ) 
	{
		$format = $test;
	}
}
print "Recognized format for BondFragment filename is: $format\n";

 
# Parse local energy and forces from the various FragmentDirs
for ($i=0;$i<=$#Factors;$i++) 
{
	my $filename = sprintf($format, $i);
	my @LocalEnergy;
	open(TE, "<$Inputdir/BondFragment".$filename."/$Prefix.energy.all") or die "could not open $Inputdir/BondFragment".$filename."/$Prefix.energy.all: $!";
# skip first line containing the header
	$line = <TE>;	
	#	($Time, $Total, $Kinetic, $NonLocal, $Correlation, $Exchange, $Pseudo, $Hartree, $Gauss, $Ewald, $ETotal)
	while ($line = <TE>) {
		chomp($line);
		my @tmp = split(/[ \t]+/, $line);
		for $k ( 0 .. $#tmp) {
			$tmp[$k] *= $Factors[$i];
		}
		push @LocalEnergy, [ @tmp ];
	}
	close(TE);

	# sum up energy
	for $j ( 0 .. $#LocalEnergy ) 
	{
		#print " @{$LocalEnergy[$j]} \n";
		if ( $i == 0) {
			push @TotalEnergy, [ @{$LocalEnergy[$j]} ];
		} else {
			for $k ( 0 .. $#{$LocalEnergy[$j]} ) {
				$TotalEnergy[$j][$k] += $LocalEnergy[$j][$k];
			}
		}
	}

	my @LocalForce;
	open(FORCE, "<$Inputdir/BondFragment".$filename."/$Prefix.forces.all") or die "could not open $Inputdir/BondFragment".$filename."/$Prefix.forces.all: $!";
# skip first line containing the header
	$line = <FORCE>;
	while (<FORCE>) {
		next if /MeanForce:/;
		chomp;
		my @tmp = split(/[ \t]/);
		push @LocalForce, [ @tmp ];	
	}
	close(FORCE);

	# sum up forces
	for $j ( 0 .. $#LocalForce) {	# atom in fragment
		# Add local forces times their factor
		for $k ( 0 .. $#{$LocalForce[$j]} ) {	# force type
			for $l ( 0 .. $#{$ForceFactors[$i]} ) {	# whole molecule atom
				#print "fragment $i\tfragment atom $j\t force type $k\t molecule atom $l\n";
				if (0) { # ($j == 0) && ($i == 0)) {
					if ($k >= 5) {		# forces start at column 5
						$TotalForce[$l][$k] = $LocalForce[$j][$k] * $ForceFactors[$i][$l][$j];
					} else {	# take correct position only if local fragment atom relates to whole molecule atom
						if ($ForceFactors[$i][$l][$j] != 0) {
							$TotalForce[$l][$k] = $LocalForce[$j][$k];
 						}
					}
				} else {
					if ($k >= 5) {    # forces start at column 5
						$TotalForce[$l][$k] += $LocalForce[$j][$k] * $ForceFactors[$i][$l][$j];
					}	else {  # take correct position only if local fragment atom relates to whole molecule atom
						if ($ForceFactors[$i][$l][$j] != 0) {
							$TotalForce[$l][$k] = $LocalForce[$j][$k];
						}
					}
				}
			}
		}
	}
}

# output summed up approximate total energy
open(TE, ">$Inputdir/$Prefix.energy.all") or die "could not open $Inputdir/$Prefix.energy.all: $!";
printf(TE "#Level\t%s\t\t%s\t\t%s\t\t%s\t%s\t%s\t%s\t\t%s\t\t%s\t\t%s\t\t%s\t\t%s\n",
      "Time",
      "Total",
      "Kinetic", "NonLocal",
      "Correlation", "Exchange",
      "Pseudo", "Hartree",
      "-Gauss",
      "Ewald",
      "IonKin",
      "ETotal");
for $j ( 0 .. $#TotalEnergy ) 
{
	print TE ($j+1)."\t";
	for $k ( 0 .. $#{$TotalEnergy[$j]} ) {
  	print TE "$TotalEnergy[$j][$k]\t";
	}
	print TE "\n";
}
print TE "\n";
close(TE);

# output summed up approximate total forces
open(FORCE, ">$Inputdir/$Prefix.forces.all") or die "could not open $Inputdir/$Prefix.forces.all: $!";
printf(FORCE "#%s\t%s\t%s\t\t%s\t\t%s\t\t%s\t\t%s\t\t%s\t\t%s\t\t%s\t\t%s\t\t%s\t\t%s\t\t%s\t\t%s\t%s\t%s\t%s\t\t%s\t\t%s\n",
	  "Type", "No",
	  "Pos0", "Pos1", "Pos2",
	  "Total0", "Total1", "Total2",
	  "Local0", "Local1", "Local2",
	  "NLocal0", "NLocal1", "NLocal2", 
    "Magnetic0", "Magnetic1", "Magnetic2", 
	  "Ewald0", "Ewald1", "Ewald2");
for $i ( 0 .. $#TotalForce ) {
	for $j ( 0 .. $#{$TotalForce[$i]} ) {
		print FORCE "$TotalForce[$i][$j]\t";
	}
	print FORCE "\n";
}
print FORCE "\n";
close(FORCE);


exit 0;