#!/usr/bin/perl -w

#################################################################################
# This script extracts time estimates obtained under the AHRS approach          #
# implemented in PAML (ver. 3.15)                                               #
#                                                                               #
# Comments to: sarisbro at uottawa dot ca                                       #
# Date: May. 2006                                                               #
#################################################################################

use strict;
use File::Copy;

my $NbMaxClusters = 10;
my $NbReps = 1000;

#################################################################################
# MAIN                                                                          #
#################################################################################

my($OStype,$pwd) = get_pwd();
my $curdir = ''; if($curdir eq ''){$curdir = $pwd;}else{$curdir = $ARGV[0];}
my $full_dir_path;

open(OUTFILE,">$curdir/Ktimes.xls") || die "\nCould not create out file: $!\n";
open(OUTTREEFILE,">$curdir/out.trees") || die "\nCould not create outtree file: $!\n";
print OUTFILE "K\tlnL";
for(my $j=41; $j<80; $j++){
	print OUTFILE "\tnode_$j\tsd_$j";
}

for(my $i=1; $i<$NbMaxClusters+1; $i++){
	if(-e "$curdir/k$i"){
		print "\n K= $i";
		$full_dir_path = "$curdir/k$i";
		get_res_with_max_lnL($full_dir_path, $i);
		extract_times($full_dir_path, $i);
		unlink "$full_dir_path/mycodeml";
		unlink "$full_dir_path/mybaseml";
	}else{
		print OUTFILE "\n$i";
	}
}

close(OUTFILE);
close(OUTTREEFILE);
print "\n\n\aAll done!\n";

#################################################################################
# SUBS                                                                          #
#################################################################################
sub get_res_with_max_lnL{ # start_ana_over_K.pl ran baseml $NbReps times. This
# copies the rep with the largest log-likelihood to $curdir/k$cluster_id
	my $path = shift @_;
	my $cluster_id = shift @_;
	
	my $best_rep = -1;
	my $best_lnL =  -1000000000;
	my @lnL_line;
	
	chdir "$path/k$cluster_id";
	for(my $i=0; $i < $NbReps; $i++){
		open(RESFILE,"$path/rep$i/clock.test_6_k$cluster_id.out") || warn "\nCould not open res file for cluster $cluster_id rep $i... $!\n";
		while(<RESFILE>){
			chomp(my $curline = $_);
			if($curline =~ /^lnL\(/){
				@lnL_line = split(/\):/,$curline);           # lnL is at $lnL_line[1]
				next;
			}
		}
		close(RESFILE);
		if($lnL_line[1] > $best_lnL){
			$best_lnL = $lnL_line[1];
			$best_rep = $i;
		}
	}
	
	print ">$best_rep<\n";
	copy("$path/rep$best_rep/clock.test_6_k$cluster_id.out","$path/clock.test_6_k$cluster_id.out");
}
#################################################################################
sub extract_times{ # and log-likelihood values for the whole tree (all the genes)
	my $path = shift @_;
	my $cluster_id = shift @_;
	
	my($step3,@lnL_line,@time_line,@time_est,$estimates,$getTrees1,$getTrees2,$brIdTree,$rateClassTree,$divTimeTree);
	
	$estimates = "";
	$getTrees1 = $getTrees2 = 0;
	$rateClassTree = "";
	open(RESFILE,"$path/clock.test_6_k$cluster_id.out") || die "\nCould not open res file for cluster $cluster_id... $!\n";
	while(<RESFILE>){
		chomp(my $curline = $_);
		$step3 = 0;
		if($curline =~ /^Step 3/){
			$step3 = 1;
		}
		elsif($curline =~ /^lnL\(/){
			@lnL_line = split(/\):/,$curline);           # lnL is at $lnL_line[1]
		}
		elsif($curline =~ /^Node/ && $curline =~ /Jeffnode/){
			@time_line = split(/Time/,$curline);
			if($curline =~ /\+\-/){
				@time_est  = split(/\+\-/,$time_line[1]);
				$time_est[0] =~ s/\s+//g;
				$time_est[1] =~ s/\s+//g;
				if($time_est[1] =~ /nan/){
					$time_est[1] = "";
				}
				$estimates .= "$time_est[0]\t$time_est[1]\t";
			}else{
				$time_line[1] =~ s/\s+//g;
				$estimates .= "$time_line[1]\t\t";
			}
			next;
		}
		# This is to get trees
		elsif($curline =~ /^rate group/){
			$getTrees1 = 1;
		}
		elsif($getTrees1 && $curline =~ /^\(/){
			$rateClassTree .= "\n$curline";
			$rateClassTree =~ s/ \#0.0000//g;
			$rateClassTree =~ s/\#/ \_/g;
			$getTrees1 = 0;
		}
		elsif($curline =~ /Tree with node ages for TreeView/){
			$getTrees2 = 1;
		}
		elsif($getTrees2 && $curline =~ /\:/){
			$divTimeTree = "$curline";
		}
		elsif($getTrees2 && $curline =~ /\d\_/){
			$brIdTree = "$curline";
		}
		elsif($curline =~ /Nodes and Times/){
			$getTrees2 = 0;
		}
	}
	close(RESFILE);
	
	print OUTFILE "\n$cluster_id\t$lnL_line[1]\t$estimates";
	if($cluster_id < 2){
		print OUTTREEFILE "$brIdTree\n$divTimeTree $rateClassTree\n";
	}else{
		print OUTTREEFILE "$divTimeTree $rateClassTree\n";
	}
}
#################################################################################
sub get_pwd{
	my $OStype = ''; my($this_pwd,@dum);
	# test if OS is Windows
	system "ver > ostype.txt";
	open(TESTFILE,"ostype.txt") || warn "\ncould not open ostype.txt(1)... $!\n";
	while(<TESTFILE>){
		chomp(my $curline = $_);
		if($curline =~ /Windows/){
			$OStype = "Windows";
			next;
		}
	}
	close TESTFILE;

	# test if OS is linux/unix, & pretty much assumes this type if you get here
	if($OStype eq ''){
		system "uname > ostype.txt";
		open(TESTFILE,"ostype.txt") || die "\ncould not open ostype.txt(L)... $!\n";
		while(<TESTFILE>){
			chomp($OStype = $_);
		}
		close TESTFILE;
	}
	unlink "ostype.txt";

	if($OStype eq "Windows"){
		system "dir > pwd.txt";
		open(PWDFILE,"pwd.txt") || die "\ncould not open pwd.txt(W)... $!\n";
		while(<PWDFILE>){
			chomp(my $curline = $_);
			if($curline =~ /Directory of/){
				@dum = split(/ of /,$curline);
				$this_pwd = $dum[1];
				next;
			}
		}
		system "cls";
	}else{
		system "pwd > pwd.txt";
		open(PWDFILE,"pwd.txt") || die "\ncould not open pwd.txt(L)... $!\n";
		while(<PWDFILE>){
			chomp($this_pwd = $_);
		}
		system "clear";
	}
	close PWDFILE;
	unlink "pwd.txt";
	$this_pwd =~ s/\\/\//g if $OStype eq "Windows";

	return($OStype,$this_pwd);
}
#################################################################################
# THAT'S ALL FOLKS                                                              #
#################################################################################
