#!/usr/bin/perl
#aln_order_v2.0.pl version 2.0.1
#   author: Sonja J. Prohaska
#institute: Bioinformatics, Institute for Computer Science
#  address: Haertelstr. 16-18, 04107 Leipzig, Germany
#     date: 04.12.2007
# language: perl

# v2.0.0: introduction of option -rr and corresponding code
#         introduction of option -io and corresponding code
#         @seqs is now loaded with sequences in input order
#         output of distance matrix turned off if distance matrix was input
#         reformating of output

# v2.0.1: bug fix: scrambled gene order when reading from alignment file 
#         bug fix: runtime error when reading scoring matrix  

use strict;

use vars qw/@seqs %distmat/;

my $usage= << "JUS";
  usage: aln_order_v2.0.pl [OPTIONS] [MULTIFASTA SCORINGMATRIX | -dist=DISTANCEMATRIX]
options: -d        debug
         -dmo      distance matrix (only)
         -msto     minimum spanning tree (only)
         -mhpo     minimum hamiltonian path (only)
         -rr=      calculate path for x *r*epetitions of *r*andom paths
         -io       calculate path for genes in *i*nput *o*rder
JUS

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

##default settings:
my $debug = 0; 
my $alnfile;
my $scorefile;
my $distancefile;
my $scorematref;
my $dmo;
my $msto;
my $mhpo;
my $rr = 0;
my $io;
#my @seqs;
#my %distmat;  
##----------

##setting parameters:
foreach my $option (@options){
  if    ($option =~ m/^-d$/)          { $debug = 1   }
  elsif ($option =~ m/^-dist=(.*)$/)  { $distancefile = $1; }
  elsif ($option =~ m/^-dmo$/)        { $dmo = 1; }
  elsif ($option =~ m/^-msto$/)       { $msto = 1; }
  elsif ($option =~ m/^-mhpo$/)       { $mhpo = 1; }
  elsif ($option =~ m/^-rr=(\d*)$/)   { $rr = $1; }
  elsif ($option =~ m/^-io$/)         { $io = 1; }
  else  { die "Invalid command line option: \'$option\'!\n$usage"; }
}

die "$usage" unless (scalar(@others)== 2 | defined($distancefile));
$alnfile=$others[0];
$scorefile=$others[1];

unless ($dmo == 1 | $msto == 1 | $mhpo == 1) {
  $dmo=$msto=$mhpo=1;
}

#(0)# read distance matrix, if available
if (defined($distancefile)) {
  my ($distmatref,$names)=read_matrix($distancefile);
  %distmat=%$distmatref;
#  @seqs=keys(%distmat);
  @seqs=@$names;
}
else {
#(1)# read alignment
  my ($aliref, $namesref)=read_multifasta($alnfile);
#  @seqs=keys(%$aliref);
  foreach my $k (sort {$a <=> $b} (keys(%{$namesref->[1]}))) {
    push @seqs, $namesref->[1]->{$k};
  }
  if ($debug) {print STDERR "sequence: @seqs\n";}
  my $alilen=scalar(@{$aliref->{$seqs[0]}});
  
#(2)# read scoring matrix
  my $null;
  ($scorematref, $null)=read_matrix($scorefile);
  if ($debug) {
    print STDERR "scoring matrix from file \'$scorefile\':\n";
    print_matrix($scorematref, *STDERR);
    print STDERR "-----\n";
  }

#(3)# calculate distance matrix
  
  for (my $i=0; $i<scalar(@seqs); $i++) {
    
    for (my $j=$i; $j<scalar(@seqs); $j++) {
      
      for (my $pos=0; $pos<$alilen; $pos++) {
	if (exists($scorematref->{$aliref->{$seqs[$i]}->[$pos]}) &
	    exists($scorematref->{$aliref->{$seqs[$i]}->[$pos]}->{$aliref->{$seqs[$j]}->[$pos]})) {
	  my $posdist = $scorematref->{$aliref->{$seqs[$i]}->[$pos]}->{$aliref->{$seqs[$j]}->[$pos]};
	  ${$distmat{$seqs[$i]}}{$seqs[$j]}+=$posdist;
	  ${$distmat{$seqs[$j]}}{$seqs[$i]}+=$posdist;
	}
	else {
	  print "Invalid character \"" . $aliref->{$seqs[$i]}->[$pos] . "\" in >" . $seqs[$i] . "< or \"" . $aliref->{$seqs[$j]}->[$pos] . "\" in >" . $seqs[$j] . "<.\n"; 
	}
      }
    }
  }
}

#(3.1)# compute list of distances
my @distlist;

for (my $i=0; $i<scalar(@seqs)-1; $i++) {
  for (my $j=$i+1; $j<scalar(@seqs); $j++) {
    push @distlist, [%distmat->{$seqs[$i]}->{$seqs[$j]}, $seqs[$i], $seqs[$j]];
  }
}

if (defined($dmo) == 1 & defined($distancefile) == 0) {
  print STDOUT "# distance matrix:\n";
  print_matrix(\%distmat, *STDOUT);
  print STDOUT "# -----\n";
}

if ($msto) {
#(4)# calculate minimum spanning tree (best solution but not a linear ordering)
    # Kruskal's algorithm (Graphs: Theory and Algorithms, K. Thulasiraman) 
  
  my @edges;    #sorted list of edges
                #[0] weigth [1] vertex_1 [2] vertex_2
  my %concoms;  #connected components
                # {seq} -> [0] label/color of component seq is part of
  my @spanning; #list of edges of the spanning tree
                #[0] weigth [1] vertex_1 [2] vertex_2
  
#(4.1)# sort edges ascending

  @edges = sort {$a->[0] <=> $b->[0]} @distlist;
  
  if ($debug) {
    print STDERR "sorted list of distances:\n";
    print_list(\@edges, *STDERR);
    print STDERR "-----\n";
  }
#(4.2)# put each vertex in its own connected component
      # by labeling each vertex with a different color

  my $l=0;
  foreach my $i (@seqs) {
    $concoms{$i}=$l;
    $l++;
  }
  
#(4.3)# for each edge e_i (in the order of minimun weight)
      # if e_i connects two components (graph will stay acyclic)
      # a) add the edge e_i to the spanning tree 
      # b) unite the connected components of the two vertices
      #    (recolor vertices)

  foreach my $e (@edges) {
    my %nr_cc;
    # if e_i connects two components (graph will stay acyclic)
    my $cc1=$concoms{$e->[1]}; #connected component 1
    my $cc2=$concoms{$e->[2]}; #connected component 2
    unless ($cc1 == $cc2) {
      # a)
      push @spanning, $e;
      # b)
      my $label; #use label with smaller number
      if ($cc1 > $cc2) {
	$label = $cc2;
      }
      else {
	$label = $cc1;
      }
      # recoloring
      
      foreach my $v (keys(%concoms)) {
	$nr_cc{%concoms->{$v}}++;
	if (%concoms->{$v} == $cc1 | %concoms->{$v} == $cc2) {
	  %concoms->{$v}=$label;
	}
      }
      
#(4.4)# once I join the last two connected components => one component. Done.
      if (scalar(keys(%nr_cc)) == 2) {
	print "# minimum spanning tree (MST):\n";
	my $len= print_list(\@spanning, *STDOUT);
	print "MST [$len]\n";
	print "# -----\n";
      }
    }
  }
}

#???# How to compute *all* optimal spanning trees?
#???# How to compute suboptimal spanning trees?

if ($mhpo) {
#(5)# calculate minimum hamiltonian path (best linear ordering)
#     use dynamic programming

# minimize over minimum hamiltonian pathes starting at each vertex
  my $min;
  my %pathes;
  
  for (my $start=0; $start<scalar(@seqs); $start++) {
    my %holdit;
    my @S;
    for (my $s=0; $s<scalar(@seqs); $s++) {
      push @S, $s unless ($s == $start);
      #print "$seqs[$s] -> $s\n";
    }
    print STDERR "Computing hamiltonian path " . ($start+1) . " of "
	. scalar(@seqs) . "...\n";
    #print "S>@S<\n";
    my ($tmp,$order)=DP_function($start,\@S, \%holdit);
    push @{%pathes->{$tmp}}, $order;
    if ($tmp < $min || defined($min)==0) { $min=$tmp;}
  }
  
  print STDERR "done.\n";
  
  print "# minimum hamiltonian path(s) (MHP):\n";
  foreach my $path (@{%pathes->{$min}}) {
    print_path($path, $min, 'MHP');
  }
  print "# -----\n";

  # calculate path lengthes of random pathes for 'rr' repetitions
  if ($rr) {
    print "# length of $rr random paths (RR):\n";
    foreach my $x ((1..$rr)) {
      # generate a random path
      my $rpath=random_path();
      # compute its length
      my $len=path_length($rpath,\%distmat);
      print_path($rpath, $len, 'MHPRR');
    }
    print "# -----\n";
  }
  # calculate path lengthes of path specified by input order
  if ($io) {
    print "# path length of genes in input order (IO):\n";
    my $iopath = [0..$#seqs];
    my $len=path_length($iopath,\%distmat);
    print_path($iopath, $len, 'IO');
    print "# -----\n";
  }
}

sub path_length {
  my ($path, $distmatref)=@_;
  my $len=0;
  for (my $i=0; $i<scalar(@$path)-1; $i++) {
    $len += $distmatref->{$seqs[$path->[$i]]}->{$seqs[$path->[$i+1]]};
  }
  return $len;
}

sub random_path {
  my @randpath;
  my @tmp = (0..$#seqs);
  for (my $i=0; $i<scalar(@seqs); $i++) {
    my $tmplen = scalar(@tmp);
    my $rand = int(rand($tmplen));
    my $select = splice(@tmp, $rand, 1);
    die "Exit: Atempt to select non-existing array element!" if $select eq '';
    push @randpath, $select;
  }
  return \@randpath;
}

sub print_path {
  my ($path, $len, $ID)=@_;
  my @tmp = @seqs[@$path[0..$#$path]];
  print "$ID @tmp [$len]\n";
}

sub DP_function {
  my ($i,$S, $holdit)=@_;
  my $min;    # length of the minimum hamiltonian path
  my @order;  # minimum hamiltonian path
  
  # c(i,S) for S={j} = dij
  if (scalar(@$S) == 1) {
    $min = %distmat->{$seqs[$i]}->{$seqs[$S->[0]]};
    push @order, $i, $S->[0];
  }
  # c(i,S) = min_{j element in S} dij + c(j, S-{j})
  else {
    foreach my $j (@$S) {
      my @S_min_j = grep { $_ != $j } @$S;
      my $t=tostring($j,\@S_min_j);
      # temporary holders for min and order
      my $t_min;
      my @t_order;
      if (exists($holdit->{$t})) {
	$t_min=%distmat->{$seqs[$i]}->{$seqs[$j]} + $holdit->{$t}->[0];
	push @t_order, $i, @{$holdit->{$t}->[1]};
      }
      else {
	# this is the recursion!
	my ($min_dp, $order_dp)=DP_function($j, \@S_min_j, $holdit);
	# this is the dynamic programming part!
	$t_min=%distmat->{$seqs[$i]}->{$seqs[$j]} + $min_dp;
	push @t_order, $i, @$order_dp;
      }
      # taking the minimum over all j
      if ($t_min < $min || defined($min)==0) { 
	$min=$t_min;
	@order=@t_order;
      }
    }
  }
  #print "i>$i< S>@$S< min>$min<\n";
  my $tt=tostring($i,$S);
  $holdit->{$tt}=[$min, \@order];
  return $min, \@order;
}

sub tostring {
  my ($i, $S) = @_;
  my @t=sort(@$S);
  my $tt=join(":",$i,@t);
  return $tt;
}

#  takes: filename of an alignment in MULTIFASTA format
#returns: ref to alignment hash, key = name, value = list of chars
#         ref to list of two hashes (namesref),
#         first  (NaNr): key = name (Na), value = number (Nr)
#         second (NrNa): key = number (Nr), value = name (Na)
#         while Nr refects the order of names in the alignment file     
sub read_multifasta {
  my $file = shift;
  my %ali;
  my %NaNr;
  my %NrNa;
  my $gap='-';
  my $index=0;
  open (MFA, "<$file") or die "Can\'t open file \'$file\'!\n";
 outer: while(<MFA>) {
   chomp;
   if ($_ =~ m/^>\s*(\S+)\s*/) {
     my $org = $1;
     unless (exists $NaNr{$org}) {
       $index++;

       $NaNr{$org}=$index;
       $NrNa{$index}=$org;
#       if ($debug) {print "<$index>-<$org>\n"}
     }
#     @{$ali{$org}} = ();                                                       
     while(<MFA>) {
       chomp;
       if ($_ =~ m/^[a-z$gap]*$/i) {
         push @{$ali{$org}}, split('',$_);
       }
       elsif ($_ =~ m/^>\s*(\S+)\s*/) {
         redo outer;
       }
       elsif ($_ =~ m/^\s*$/) {
         next;
       }
       else {
         print STDERR "Warning: Something wrong in multifasta \'$file\':\n$_\n";
         last;
       }
     }
   }
 }
  close MFA;
  return \%ali, [ \%NaNr, \%NrNa ];
}



#  takes: filename of an MATRIXFILE
#returns: ref to matrix hash, key1 = line_name, key2 = column_name
#         (ganz nach dem Grundsatz *Z*eilen *z*uerst, *S*palten *s*paeter)
#         value = value at key1,key2
#         and input order of column_names
sub read_matrix {
  my $mfile=shift;
  my %matrix;
  my @index;
  open (IN, "<$mfile") or die "Can\'t read/open \'$mfile\'!\n";
  while (<IN>) {
    if ($_ =~ m/^\s*$/) { next }
    chomp;
    my @line = split;
    my $first = shift @line;
    if ($. == 1) {
      foreach my $i (@line) {
	push @index, $i;
      }
    }
    else {
      for (my $j=0; $j<scalar(@index); $j++) {
	$matrix{$first}->{$index[$j]}=$line[$j];
      }
    }
  }
  return \%matrix, \@index;
}

sub print_matrix {
  my ($hash,$to)=@_;
  foreach my $key1 (keys(%$hash)) {
    print $to "$key1";
    foreach my $key2 (keys(%{$hash->{$key1}})) {
      print $to "\t" . $hash->{$key1}->{$key2};
    }
    print $to "\n";
  }
}

sub print_list {
  my ($list, $to)=@_;
  my $mstlen=0;
  foreach my $i (@$list) {
    print $to "# $i->[1] $i->[2] [$i->[0]]\n";
    $mstlen += $i->[0];
  }
  return $mstlen;
}
