#!/usr/bin/perl -w
#
# author: axel wintsche
# mail: wint(spider monkey)bioinf(dot)uni-leipzig(dot)de

use strict;

# coloumn from which stats where choosen
# 1: number of genes in a n organism
# 2: number of domains in an organism
# 3: average number of domains per gene in an organism
my $statcol = 1; 

# file with information on number of known organisms and explicit parent
#my $info_file = "/homes/martini/praktikum/prak09/viechcount";


$#ARGV == 1 or die "try addStats.pl <statlist> <newick tree file>\n";



open(STAT,"<$ARGV[0]") or die "can't find statlist\n";

my %stats;

foreach(<STAT>) {
	
	#use first col (species name) as key and $statcol as value
	#pay attention that the species name is the same as used in the tree file!
	
	my @temp = split(/\t/, $_);
	chomp @temp;
	
	$temp[0] =~ s/[\(\),:]//g;
	
	#due to inconsistent files man has to check if a key already exists
	if( exists $stats{$temp[0]} ) {
		#use the maximal value, also other values e.g. avg could be used here
		if( $stats{$temp[0]} < $temp[$statcol] ) { $stats{$temp[0]} = $temp[$statcol]; }
	}
	else {
		$stats{$temp[0]} = $temp[$statcol];
	}

}

close(STAT); 


open(TREE, "<$ARGV[1]") or die "can't find tree file\n";

my $tree = "";

foreach(<TREE>) {

	$tree .= $_;
	chomp $tree;
}

close(TREE);

my $tree_with_stats = "";
my $species = "";
my $name_contain_bracket = 0;
my @stack;

for(my $i = 0; $i < length $tree; $i++) {
	my $char = substr($tree, $i, 1);
	
	if($char eq "(") {
				
		$tree_with_stats .= $char;
		#create a new element in stack
		push( @stack, ["x"]);
	}
	elsif($char eq ")") {
	
		$tree_with_stats .= $char."[";
		# calculating avg and print stats
		my $element = pop @stack;
		#add mimimum:
		$tree_with_stats .= (@$element[0])."/";
		#add average:
		die "$tree_with_stats" if(!defined @$element[3]);
		$tree_with_stats .= sprintf "%.2f/", (@$element[2]/@$element[3]);
		#add maximum:
		$tree_with_stats .= (@$element[1])."/";
		#add #organism in subtree
		$tree_with_stats .= (@$element[3]);
		
		join_with_ancestor($element);
		
		$tree_with_stats .= "]";
		$species = "";	#set to empty string
	}
	elsif($char eq ",") {
		$tree_with_stats .= $char;
		#it follows the next tree node, so clear species name
		$species = "";
	}
	elsif($char eq ":") {
		#species name is complete,
		#if leaf node (species name is not the empty string)
		#get genecount and add to stackentry
		if( $species ne "" ) {
			
			if(exists $stats{$species}) {
				add_gene_count($stats{$species});
				$tree_with_stats .= " [".$stats{$species}."]";
			}
			else { print STDERR "$species\n"; }
			
		} 
		$tree_with_stats .= $char;
	}
	else {
		$species .= $char;
		$tree_with_stats .= $char;
	}

}

add_outgroup();

print "$tree_with_stats\n";

##TODO: add dump node as new "root" to show stats also for whole tree...
sub add_outgroup {
	chop $tree_with_stats;
	$tree_with_stats = "(".$tree_with_stats.":1,root[min/avg/max/#organism in subtree]:1);";
}

sub add_gene_count {
	#if this is the first call (@stack is empty)
	#create new entry,
	#else alter the element on top of the stack
	
	my $element = pop @stack;
	my $value = shift;	
	
	if(@$element[0] eq "x") { push(@stack, [$value,$value,$value,1]); } #elemtent(min, max, sum, #child nodes)
	
	else {

		if( $value < @$element[0] ) { @$element[0] = $value; }
		elsif( $value > @$element[1] ) { @$element[1] = $value; }
		
		@$element[2] += $value;	#add value to the sum
		@$element[3]++;		#increase number of child nodes
		
		push(@stack, $element);
		
	}
}

sub join_with_ancestor {
	
	if($#stack == -1) { return; }
	
	my $ancestor = pop @stack;
	
	#check if ancestor is uninitialised yet
	if( @$ancestor[0] eq "x" ) {
		push( @stack, shift ); #take current stats as ancestor stats and leave
		return;
	}
	
	my $current = shift;
	
	if( @$current[0] < @$ancestor[0] ) { @$ancestor[0] = @$current[0]; }
	if( @$current[1] > @$ancestor[1] ) { @$ancestor[1] = @$current[1]; }
	
	@$ancestor[2] += @$current[2];
	@$ancestor[3] += @$current[3];
	
	push( @stack, $ancestor)
}
