#!/usr/bin/perl
#WagnerTest.pl version 1.0
#   author: Sonja J. Prohaska
#institute: Bioinformatics, Institute for Computer Science
#  address: Haertelstr. 16-18, 04107 Leipzig, Germany
#     date: 23.03.2010

use strict;

my $usage= << "JUS";
  usage: WagnerTest.pl [OPTIONS] ALIGNMENTFILE TREEFILE
options: -d          debug
         -cw         default: alntype is clustalw
         -mfa        alntype set from clustalw to multifasta
         -lc=ignore  default: treat lower case letters as unaligned (dialign2)
         -lc=uc      treat lower case letters like upper case letters
         -r=         number of bootstrap replicates, default: 0         
                     (each bootstrap step will take as long as the actual run)
         -bootsize=  default: double the mean block length
JUS

my @options = grep {substr($_, 0, 1) eq '-'} @ARGV;
my @others  = grep {substr($_, 0, 1) =~ /^[^-]/} @ARGV;

##default settings:
my $debug   = 0 ;         # no debugging
my $alntype = 0 ;         # default: 0 = cw, 1 = mfa
my $lc      = 0 ;         # default: 0 = lc are ignored, 1 = lc are as UC 
my $r       = 0 ;         # default: 0
my $p       = 0 ;         # default: 0 (do not overrule estimation of
                          #            the ideal block length for sampling)
##----------

##setting parameters:
foreach my $option (@options){
  if    ($option eq "-d")          { $debug = 1; }
  elsif ($option eq "-cw")         { $alntype = 0; }
  elsif ($option eq "-mfa")        { $alntype = 1; }
  elsif ($option eq "-lc=ignore")  { $lc = 0; }
  elsif ($option eq "-lc=uc")      { $lc = 1; }
  elsif ($option =~ m/^-r=(\d+)$/) { $r = $1; }
  elsif ($option =~ m/^-bootsize=(\d+)$/) { $p = $1 }
  else  { die "Invalid command line option: \'$option\'!\n$usage"; }
}

die "Two files expected: ALIGNMENTFILE, TREEFILE!\n$usage" 
    unless scalar(@others)==2;
foreach my $file (@others) {
  die "'$file' no such file or directory!" unless -e $file;
}
my $alnfile =$others[0];
my $treefile=$others[1];

#read tree:
my ($spsref, $lcamref, $dtmref) = read_tree($treefile);

if ($debug) {
    print "\nlca matrix:";
    print_matrix($lcamref, $spsref);
    print "\ndiv time matrix:";
    print_matrix($dtmref, $spsref);
}

#read alignment:
my $aliref;
if ($alntype) {
  ($aliref)= read_mfa($alnfile);
}
else {
  ($aliref)= read_cw($alnfile);
}

if ($debug) {
    print "alignment:\n";
    print_ali($aliref);
}

my @tmp=sort(@$spsref);
print "# species in      tree: @tmp\n";
my @tmp=sort(keys(%$aliref));
print "# species in alignment: @tmp\n";

#compute valid quartets:
my $quatref = valid_quartets($spsref, $lcamref, $dtmref);
#print_quartets($quatref);

#compute relative rates:
my ($rrtref, $bl) = relrat($quatref, $aliref, $lc);
print_rrtres($rrtref, '1');
print "#--------------------\n";

#bootstrap sequences:
if ($r) {
  for (my $i=1; $i<=$r; $i++) {
    my $bootaliref=random_alignment($aliref, $bl, $p);
    print STDERR "replicate $i of $r:\n";
    my ($bootrrtref, $bbl) = relrat($quatref, $bootaliref, $lc);
    print "# replicate $i\n";
    print_rrtres($bootrrtref, '0');
  }
}

sub relrat {
    my ($quatref, $aliref, $lc) = @_;
    my $gap = '-';
    my @rrtres;
    my $iq =0;
    my $bl;
    foreach my $q (@$quatref) {
        $iq++;
        print STDERR "$iq of " . scalar(@$quatref) . " quartets...\n";
	my $cXO;
	my $cXA;
	my $cXB;
	my $cAB;
	my $u; 
	my $block_l = 0;  #length of a contiguous block where kO eq kX 
	my @block_l;      #collects block lengthes
	my $seqO = $aliref->{$q->[0]} if exists($aliref->{$q->[0]});
	my $seqX = $aliref->{$q->[1]} if exists($aliref->{$q->[1]});
	my $seqA = $aliref->{$q->[2]} if exists($aliref->{$q->[2]});
	my $seqB = $aliref->{$q->[3]} if exists($aliref->{$q->[3]});
	my $T2   = $q->[4];
	unless ($seqO & $seqX & $seqA & $seqB) {
	    print "quartet @$q not computed!\n";
	    next;
	}
	my $l = length($seqO);
	
	for (my $i=0; $i<$l; $i++) {
	    my $kO = substr($seqO, $i, 1);
	    my $kX = substr($seqX, $i, 1);
	    my $kA = substr($seqA, $i, 1);
	    my $kB = substr($seqB, $i, 1);
	    #is it correct to replace lc letters (dialign) by gaps?
	    #we can do it, because we can alsways lengthen the alignment
	    #to get a block composed of lc letters an gaps and
	    #a block where lc letters are replaced by gaps:

	    # Aa ATGGCC     AT -- GGCC 
	    # Bb AaccCC --> A- cc --CC
	    # Cc ATttCC     AT tt --CC

	    foreach my $k ($kO, $kX, $kA, $kB) {
		if ($k =~ m/[a-z]/ ) {
		    if ($lc == 0) {
			$k = $gap;
		    }
		    elsif ($lc == 1) {
			$k = uc($k);
		    }
		}
	    }
	    if ($kO eq $gap) { next }
	    if ($kO eq $kX || $kO eq $kA || $kO eq $kB) { $u++ }
	    else {
		my $g = 1;
		if ($kX eq $gap) { $g *= 0.75 }
		if ($kA eq $gap) { $g *= 0.75 }
		if ($kB eq $gap) { $g *= 0.75 }
		$u += (1-$g);
	    }
	}
	
	for (my $i=0; $i<$l; $i++) {
	    my $kO = substr($seqO, $i, 1);
	    my $kX = substr($seqX, $i, 1);
	    my $kA = substr($seqA, $i, 1);
	    my $kB = substr($seqB, $i, 1);
	    foreach my $k ($kO, $kX, $kA, $kB) {
		if ($k =~ m/[a-z]/ ) {
		    if ($lc == 0) {
			$k = $gap;
		    }
		    elsif ($lc == 1) {
			$k = uc($k);
		    }
		}
	    }
	    if ($kO eq $gap || $kX eq $gap) { 
	      #corrected version, block length distribution is now correct!
	      unless ($block_l == 0) {
		push @block_l, $block_l;
		$block_l = 0;
	      }
	      next;
	    }
	    if ($kX ne $kO) {
	      unless ($block_l == 0) {
		push @block_l, $block_l;
		$block_l = 0;
	      }
	      next;
	    }
	    $cXO++;
	    $block_l++;
	    if ($kX eq $kA) { $cXA++ }
	    if ($kX eq $kB) { $cXB++ }
	    if ($kX eq $kA && $kX eq $kB) { $cAB++ }
	    if ($kA eq $gap) {
		$cXA += 0.25;
		if ($kB eq $kX) { $cAB += 0.25 }
	    }
	    if ($kB eq $gap) {
		$cXB += 0.25;
		if ($kA eq $kX) { $cAB += 0.25 }
	    }
	    if ($kA eq $gap && $kB eq $gap) { $cAB += 0.0625 }
	}

        # for distribution of lengthes
	my %block_d;  #distribution
	my $sum;      #temp variable
	my $mean;     #mean length
	my $s;        #standard deviation
	foreach my $j (@block_l) {
	    $block_d{$j}++;
	}
	foreach my $z (@block_l) {
	    $sum += $z;
	}
	$mean = $sum/(scalar(@block_l));
	$sum = 0;
	foreach my $z (@block_l) {   
	    $sum += ($z - $mean)*($z - $mean);
	}
	$s = sqrt($sum/((scalar(@block_l))-1));

	if ($debug) {
	    print "#length distribution: $q->[0], $q->[1], $q->[2], $q->[3]\n";
	    foreach my $key  (keys(%block_d)) {
		print "$key " . $block_d{$key} . "\n";
	    }
	}
	# END (for distribution of lengthes
	
	die "Something wrong!\n" if ($cXA < $cAB || $cXB < $cAB);
	
	my $cAvB = $cXA + $cXB - $cAB;
	my $omega = $cXA*(($u+$cAvB)/($cXA+$cAvB))*(($u+$cAvB)/($cXA+$cAvB));
	my $var = 2* sqrt($omega * $cXA) * (1-$cAB/$cXB) ;
	my $la2T2 = -log($cAB/$cXA);
	my $laT2 = -log($cAB/$cXB); 
	my $la2;
	my $la;
	my $z;
	my $z_prime;
	my $violated = '';

	if ($T2) {
	    $la2 = $la2T2/$T2*1000;
	    $la = $laT2/$T2*1000;
	}

	if ($var) { 
	    $z = ($cXA-$cXB)/sqrt($var);
	    $z_prime = $z/(sqrt($mean));
	}

	if ($omega/$u < 1) {
	     $violated = "violated!"; #constraint violated
	}

	push @rrtres, [$q->[0], $q->[1], $q->[2], $q->[3], $T2, $cXA, $cXB, $cAB, $u, $omega, $var, $mean, $s, $z, $z_prime, $la2T2, $laT2, $la2, $la, $violated ];
#	printf("%5s %5s %5s %5s %8.2f %8.2f %8.2f %6.3f   >>   %6.3f %6.3f\n",
#	       $q->[0], $q->[1], $q->[2], $q->[3], $cXA, $cXB, $cAB, $mean, $z, $z_prime );
	
	$bl = $mean;
    }
    return \@rrtres, $bl;
}

sub print_rrtres {
    my ($rrtref, $h) = @_;
  # $O, $X, $A, $B, $T2, $cXA, $cXB, $cAB, $u, $omega, $var, $mean, $s, $z, $z_prime, $la2T2, $laT2, $la2, $la, $violated
    if ($h) {
      printf "#%12s %12s %12s %12s %12s  c(XA)    c(XB)    c(AB)       u        q     var   block_len +/-  s            z       z'          la2T2      laT2       la2        la\n", 'O', 'X', 'A', 'B', 'T2';
    }
    foreach my $q (@$rrtref) {
	printf " %12s %12s %12s %12s %12s %8.2f %8.2f %8.2f %8.2f %8.2f %7.2f   %6.3f +/- %6.3f  >> % 7.3f % 7.3f  <<  % 10.6f % 10.6f % 10.6f % 10.6f %8s\n", @$q;
    }
}

sub print_quartets {
    my $quatref = shift;
    printf "%12s %12s %12s %12s %12s\n", 'O', 'X', 'A', 'B', 'T2';
    foreach my $q (@$quatref) {
	printf "%12s %12s %12s %12s %12s\n", @$q;
    }
}

sub print_ali {
  my $aliref = shift;
  my $sig=3;
  foreach my $key (keys(%$aliref)) {
#    print "{$key} -> >$aliref->{$key}<\n";
    print "{$key} -> >" . length($aliref->{$key}) . "< ";
    print "(first,last)=";
    print "(" . substr($aliref->{$key},     0, $sig) . ",";
    print ""  . substr($aliref->{$key}, -$sig, $sig) . ")";
    print "\n";
  }
}

#  takes: list of species (sps), lca matrix (lcam), div. time matrix (dtm)
#returns: list of valid quartets ([0]=O, [1]=X, [2]=A, [3]=B, [4]=div.time) 
sub valid_quartets {
  my ($spsref, $lcamref, $dtmref) = @_;
  my @quartets;
  #quartets (O, X, A, B) to compute:
  my ($R,$Q,$P);
  foreach my $A (@$spsref) {
    foreach my $B (@$spsref) {
      $P=lca($lcamref, $spsref, $A, $B);
      next unless ($P);
      my $div=lca_time($dtmref, $spsref, $A, $B);
      foreach my $X (@$spsref) {
	$Q=lca($lcamref, $spsref, $X, $A);
	next unless ($Q);
	next unless ($Q eq lca($lcamref, $spsref, $X, $B));
	foreach my $O (@$spsref) {
	  $R=lca($lcamref, $spsref, $O, $X);
	  next unless ($R);
	  next unless ($R eq lca($lcamref, $spsref, $O, $A));
	  next unless ($R eq lca($lcamref, $spsref, $O, $B));
	  #print "O=$O X=$X A=$A B=$B T2=$div\n";
	  push @quartets, [$O, $X, $A, $B, $div];
	}
      }
    }
  }
#  @quartets = sort sortQuartetList @quartets;
#  print join('|', @$_), "\n" for @quartets;
#  die;
  return \@quartets;
}

#tut nix gscheites
sub sortQuartetList {
  $a->[0] cmp $b->[0]
      || $a->[1] cmp $b->[1]
#      || $a->[2] cmp $b->[2]
      || bla($a, $b) #!($a->[2] eq $b->[3] && $a->[3] eq $b->[2])
}

#tut nix gscheites
sub bla {
  my ($a, $b) = @_;
  my @d = ($a->[2], $a->[3]);
  my @e = ($b->[2], $b->[3]);
  @d = ($d[0] gt $d[1]) ? $d[1,0] : $d[0,1];
  @e = ($e[0] gt $e[1]) ? $e[1,0] : $e[0,1];
  return 0 if $a->[2] eq $b->[3] && $a->[3] eq $b->[2];
  join('', @d) cmp join('', @e); 
}

#  takes: filename
#returns: list of organsisms,
#         ?
sub read_alignment {

#calls either read_cw or read_mfa
#or returns error massege

}

#  takes: filename
#returns: alignment hash (key=sp value=sequencestring)
#this method is not save! Blocks can be left out.
sub read_cw {
  my $file = shift;
  my %ali;
  my $gap='-';
  open (CW, "<$file") or die "Can\'t open file \'$file\'! $usage";
  while(<CW>) {
    chomp;
    my @line=split;
    if (@line==2) {
      if ($line[1] =~ m/^[a-z$gap]*$/i) {
	$ali{$line[0]} .= $line[1];
      }
    }
  }
  close CW;
  my $alilen=0;
  foreach my $org (keys(%ali)) {
    #print "species=>$org<\n";
    if ($alilen==0) {
      $alilen=length($ali{$org});
    }
    else {
      unless ($alilen == length($ali{$org})) {
	die "Error at alignment! Sequences of different length, characters other than [a-z], [A-Z] and '$gap' for gaps.\n";
      }
    }
  }
  return \%ali;
}

#  takes: filename
#returns: ?
#returns: alignment hash (key=sp value=sequencestring)
sub read_mfa {
  my $file = shift;
  my %ali;
  my $gap='-';
  open (MFA, "<$file") or die "Can\'t open file \'$file\'! $usage";
 outer: while(<MFA>) {
   chomp;
   if ($_ =~ m/^>\s*(\S+)\s*/) {
     my $org = $1;
     $ali{$org} = '';
     while(<MFA>) {
       chomp;
       if ($_ =~ m/^[a-z$gap]*$/i) {
	 $ali{$org} .= $_;
       }
       elsif ($_ =~ m/^>\s*(\S+)\s*/) {
	 redo outer;
       }
       else {
	 last;
       }
     }
   }
 }
  close MFA;
  return \%ali;
}

#  takes: filename
#returns: list of species (sps),
#         last common ancestor matrix (lcam), divergence time matrix (dtm)
sub read_tree {
  my $treefile = shift;

  open (TREE, "<$treefile"); 
  my $string=<TREE>; chomp $string;
  #check if >string< has the potential to be a valid tree
  #do better than...
  unless ($string =~ m/\(\S*\(\S+\)\S*\)\S*\)/) {
    die "Unexpected format - tree >$string<\n";
  }
  my $rk=0;      #runde Klammer
  my @string = split('',$string);
  my @species;   #holds species names (do not sort!)
  my @stack;     #holds references to known nodes in the tree
  my $expect=-1; #set to pos in the string, where div time is expected
  my %sphash;    #assigns numbers (pos in @species) to sps
  my @lcamatrix; #last common ancestor matrix
  my @dtmatrix;  #divergence time matrix
  for (my $i=0; $i<scalar(@string); $i++) {
      die "Pseudonode at $i - tree >$string<\n" if ($rk < 0);
      unless ($expect == -1) {
	  if ($expect < $i) {
	      #there was no div time for the last lca introduced
	      $expect = -1;
	  }
      }
      if ($string[$i] eq '(') {
	  $rk++;
      }
      elsif ($string[$i] eq ')') {
	  $rk--;
	  $expect=$i+1; #div time expected following ')'
	  #importent things, works just for binary trees:
	  #take the last two childs from the stack
	  #and replace them by there lca (an arrayref to the two childs)
	  my $left = pop(@stack);       #left child
	  die "Too many '()' at $i - tree >$string<\n" unless ($left);
	  my $right = pop(@stack);      #right child
	  die "Too many '()' at $i - tree >$string<\n" unless ($right);
	  push @stack, [$left, $right]; #parant (lca of left and right child)
	  #get all leaves of the two subtrees
	  my $leftset =_ref2sps($left);
	  my $rightset=_ref2sps($right);
	  #all pairs of leaves (one from the left, one from the right)
	  #have the same lca (namely: $stack[$#stack])
	  foreach my $lsp (@$leftset) {
	      foreach my $rsp (@$rightset) {
		  $lcamatrix[$sphash{$lsp}][$sphash{$rsp}]= $stack[$#stack];
	      }
	  }
      }
      elsif ($string[$i] eq '[') {
	  my ($time, $j) = _time(\@string, $i);
	  die "Unclosed '[' at $i - tree >$string<\n" if ($j == $i);

	  if ($expect == $i) { #div time expected
	      die "Programming Error - badness: VERY bad.\n"
		  unless (ref($stack[$#stack]) eq 'ARRAY');
	      my ($left, $right) = @{$stack[$#stack]};
	      #get all leaves of the two subtrees
	      my $leftset =_ref2sps($left);
	      my $rightset=_ref2sps($right);
	      #all pairs of leaves (one from the left, one from the right)
	      #have the same lca (namely: $stack[$#stack])
	      #and the same div time
	      foreach my $lsp (@$leftset) {
		  foreach my $rsp (@$rightset) {
		      $dtmatrix[$sphash{$lsp}][$sphash{$rsp}]= $time;
		  }
	      }
	  }
	  elsif ($expect == -1) { #div time must refere to a leave
	      die "No ancestral species(leaves) excepted at $i - tree >$string<\n" if $time;
	  }
	  else {
	      die "Wrong position of '[' at $i - tree >$string<\n";
	  }
	  $i=$j;	  
	  $expect=-1;
      }
      else {
	  if ($string[$i-1] eq '(' | $string[$i-1] eq ',') {
	      my ($sp, $j) = _sp(\@string, $i);
	      die "Something wrong at $i - tree >$string<\n" if ($j == $i-1);
	      $i=$j;
	      push @species, $sp;
	      $sphash{$sp}=(scalar(@species)-1);
	      push @stack, \$species[$#species];
          }
          else {
	      #branch length and ':'
	  }
      }
  }
  die "Too many '(' - tree >$string<\n" if ($rk > 0);
  die "Too many ')' - tree >$string<\n" if ($rk < 0);
  die "Unconected subtrees - tree >$string<\n" if (scalar(@stack) > 1); 

  if ($debug) {
      foreach my $key (keys(%sphash)) {
	  print "$key -> $sphash{$key}\n";
      }
      my $sp = _ref2sps(\@stack);
      #print "sub>@$sp<\n";
  }
  
  return \@species, \@lcamatrix, \@dtmatrix
}

sub print_matrix {
    my ($matrixref, $speciesref) = @_;
    print "matrixdim:" . scalar(@$speciesref) . "\n";
    for (my $i=0; $i<(scalar(@$speciesref)); $i++) {
	if ($i == 0) {
	    foreach my $sp (@$speciesref) {
		printf "%20s ", $sp;
	    }
	    print "\n";
	}
	for (my $j=0; $j<(scalar(@$speciesref)); $j++) {
	    if ($matrixref->[$i][$j] eq "") {
		if ($i == $j) {
		    printf "%20s ", '0';
		}
		else {
		    printf "%20s ", "--";
		}
	    }
	    else {
		printf "%20s ", $matrixref->[$i][$j]; 
	    }
	}
	print "\n";
    }
}

#private: called by read_tree and itself
#  takes: variable
#         follows nested references
#returns: all values
sub _ref2sps {
    my $ref = shift;
    my @sps;
#    print "[>" . ref($ref) . "<]\n";
    if (ref($ref) eq '') {
#	print "VALUE\n";
	push @sps, $ref;
    }
    elsif (ref($ref) eq 'SCALAR') {
#	print "SCALAR\n";
	push @sps, $$ref;
    }
    elsif (ref($ref) eq 'ARRAY') {
#	print "ARRAY ->\n";
	foreach my $i (@$ref) {
	    my $e = _ref2sps($i);
	    push @sps, @$e;
	}
#	print "<- (ARRAY)\n"
    }
    elsif (ref($ref) eq 'REF') {
#	print "REF ->\n";
	my $e = _ref2sps($$ref);
	push @sps, @$e;
#	print "<- (REF)\n";
    }
    else {
	die "Error - wrong type of EXPR: >" . ref($ref) . "<\n";
    }
    return \@sps;
}

#private: called by read_tree
sub _time {
    my ($strref, $i) = @_;
    my $time;
    for (my $j=$i+1; $j<scalar(@$strref); $j++) {
	if ($strref->[$j] eq ']') {
	    $i=$j;
	    last;
	}
	else {
	    $time .= $strref->[$j];
	}
    }
    return $time, $i;
}

#private: called by read_tree
sub _sp {
    my ($strref, $i) = @_;
#    print "sub _sp: i=$i\n";
    my $sp;
    for (my $j=$i; $j<scalar(@$strref); $j++) {
	if ($strref->[$j] eq '[' | 
	    $strref->[$j] eq ':' |
	    $strref->[$j] eq ',' |
	    $strref->[$j] eq ')' ) {
	    $i=$j-1;
	    last;
	}
	else {
	    $sp .= $strref->[$j];
	}
    }
    return $sp, $i;
}

#  takes: last common ancestor matrix (lcam), list of species (sps), sp1, sp2
#returns: lca
sub lca {
    my ($lcam, $spsref, $sp1, $sp2) = @_;
    my %tmp;
    my $lca;
    if ($sp1 eq $sp2) {
    }
    else {
	for (my $i=0; $i<scalar(@$spsref); $i++) {
	    $tmp{$spsref->[$i]}=$i;
	}
	$lca = $lcam->[$tmp{$sp1}][$tmp{$sp2}];
	unless ($lca) {
	    $lca = $lcam->[$tmp{$sp2}][$tmp{$sp1}];
	}
	unless ($lca) {
	    die "Last common anestor of >$sp1< and >$sp2< not known!\n";
	}
    }
    return $lca;
}

#  takes: divergence time matrix (dtm), list of species (sps), sp1, sp2
#returns: divergence time of sp1 and sp2
sub lca_time {
  my ($dtm, $spsref, $sp1, $sp2) = @_;
  my %tmp;
  my $dt;
  if ($sp1 eq $sp2) {
  }
  else {
    for (my $i=0; $i<scalar(@$spsref); $i++) {
      $tmp{$spsref->[$i]}=$i;
    }
    $dt = $dtm->[$tmp{$sp1}][$tmp{$sp2}];
    unless ($dt) {
      $dt = $dtm->[$tmp{$sp2}][$tmp{$sp1}];
    }
  }
  return $dt;
}

sub random_alignment {
  my ($aliref, $bl, $p) = @_;
  my @seqs=keys(%$aliref);
  my $N = length($aliref->{$seqs[0]});
  if (defined($p) && $p>0) {
    # use given block length $p instead of calculated block length $bl
  }
  else {
    $p= 2*$bl; #take double the block length so that there is
               #at least one complete block per sampled piece
  }

  my %bootali;
  for (my $bootN=0; $bootN<=$N-$p; $bootN+=$p){
    my $dom = (rand $N)%$N;
    foreach my $k (@seqs) {
      if ($dom <= $N-$p) {
	$bootali{$k} .= substr($aliref->{$k},$dom, $p);
      }
      else {
	$bootali{$k} .= substr($aliref->{$k},$dom, $N-$dom);
	$bootali{$k} .= substr($aliref->{$k},0, $p-($N-$dom));
      }
    }  
  }
  return \%bootali;
}
