#!/usr/bin/perl -w
## This reads the concatenated file created by jarzi.f
## and generates avg + Std. Dev. using combinations of the results
##
## VARIABLES:
##  M = Number of combinations
##  N = size of subsets

if (scalar(@ARGV) != 3) {
        printf STDERR "\n This script reads the concatenated file created by jarzi.f, \n";
        printf STDERR " generating Avg + St.dev free energy profiles, from M combinations of N members.\n";
        printf STDERR " Usage:\n";
        printf STDERR " $0 j_contactenated_out.txt M N\n\n";
        exit;
}
my $IFILE = $ARGV[0];
my $M = $ARGV[1];
my $N = $ARGV[2];

my %VAR; # Hash will contain the value of the collective variable
my %W;	 # Hash will contain an array with the work values
##################################
# First read the concatenated file
##################################
open(F0,"$IFILE") || die "Cannot open file: $!";
##  Expected format: F8.3
my $l = 8;
my $nlines = 0;
while (my $line = <F0>) {
	chomp $line;
	$nlines++;
	my @tmp = &my_split($line,$l);
	my $var = shift @tmp;
	$VAR{$nlines} = $var;
	$W{$nlines} = [@tmp];
}

####################################################
# Now check consistency in the number of simulations
####################################################
my $Ntot = scalar @{$W{1}};
printf STDERR " Number of simulations: %i\n",$Ntot;
if ($N > $Ntot) {
	printf STDERR " Error: size of subset (%i) > number of simulations (%i)\n",$N,$Ntot;
	exit;
}
for my $i (2..$nlines){
	my $tmp = scalar @{$W{$i}};
	if ($tmp != $Ntot) {
		printf STDERR " Error: line %i contains %i entries (%i expected). Check input file.\n",$nlines,$tmp,$Ntot;
		exit;
	}
}

#############################
# Create M lists of N members
#############################
my %LIST;
for my $i (1..$M) {
# Generate Ntot random numbers
	my %RAND;
	for my $j (0..$Ntot-1) {$RAND{$j} = rand()}
# sort the hash by value, keep the first N
	my $n = 0;
	my @list;
	foreach my $k (sort {$RAND{$a} <=> $RAND{$b}} keys %RAND) {
		if ($n < $N) {push @list, $k;}
		$n++;
	}
	@list = sort {$a <=> $b} @list;
	$LIST{$i} = [@list];
}

# Calculate the Boltzmann-averaged work for each M set of N members
# Calculate the average and Std. dev. of these M values
my $KB=300*1.38065E-23*6.022142E23/4.184/1000.0;
#printf "kb= %f\n",$KB;

open(Fweight,">weights.txt") || die "Cannot open file: $!";
# loop over data points
for my $i (1..$nlines) {

# loop over M sets
	my @WORK = ();
	my %WEIGTH = ();
	for my $j (1..$M) {
		my $expw = 0;
# loop over N members
		my @F = ();
		for my $k (@{$LIST{$j}}) {
			my $w = $W{$i}->[$k];
			$F[$k] = exp(-1*$w/$KB);
			$expw += $F[$k];
		}
# weigths (normalized contribution of each element to work):
		for my $k (@{$LIST{$j}}) {
			push @{$WEIGHT{$k}}, $F[$k]/$expw;
		}
# get work
		$expw = $expw / $N;
		push @WORK, -1*$KB*log($expw);
	}
	printf STDOUT "%8.3f\t%8.3f\t%8.3f\n",$VAR{$i},&avg_stdev(\@WORK);
# print weigths (normalized by N/Ntot)
	my $factor = $N/$Ntot;
	for my $nn (0..$Ntot-1) {
		printf Fweight " %.6f",&avg($WEIGHT{$nn})*$factor;
	}
	printf Fweight "\n";
}
sub my_split{
	my $line = shift @_;
	my $block = shift @_;
	my $nchar = length $line;
	if ($nchar % $block != 0) {
		printf STDERR "Number of characters in line (%i) is not a multiple of %i\n",$nchar,$block;
		exit;
	}
	my $n = $nchar / $block;
	my @list = ();
	for my $i (1..$n) {
		my $end = ($i * $block) - 1;
		my $ini = $end - $block + 1;
		my $v = substr($line, $ini, $block);
		push @list, $v;
	}
	return @list;
}

sub avg_stdev{
	my $array_ref = $_[0];
	my $sum = 0;
	my $sq_sum = 0;
	my $count = 0;
	for my $a (@{$array_ref}) {
        	$sum += $a;
        	$sq_sum += ($a*$a);
        	$count++;
	}
	my $aver = $sum/$count;
	my $stdev = (($count*$sq_sum) - ($sum*$sum))/($count*$count);
	return $aver,$stdev;
}
sub avg{
	my $array_ref = $_[0];
	my $sum = 0;
	my $count = 0;
	for my $a (@{$array_ref}) {
        	$sum += $a;
        	$count++;
	}
	my $aver = $sum/$count;
	return $aver;
}