#!/usr/bin/perl -w

#dynamic programming sequence aligner
#liana lareau
#mcb246 spring 2003

#usage: aligner.pl -m <matrixfile> -p <gap penalty (positive)> -a [global | local | noend] -o <outputfile> <seq1 (fasta)> <seq2 (fasta)>

#prints the alignment and score to stdout, and the score and traceback
#matrices to files.

use strict;

my $matrixfile;                              #input matrix - blosum format
my $global;                                  #global, local, no end gap penalty
my $gap;                                     #POSITIVE gap penalty
my $outfile;                                 #file for score + trace matrices

while (@ARGV and $ARGV[0] =~ /^-/)           #commandline options after flags
{
    $_ = shift;
    if (/^-m/) {$matrixfile = shift; next;}        
    if (/^-p/) {$gap = shift; next;}  
    if (/^-o/) {$outfile = shift; next;}
    if (/^-a/) 
    {
	my $tmp = shift; 
	if ($tmp eq 'global') { $global = 1; } 
	elsif ($tmp eq 'noend') { $global = 2; }
	else { $global = 0; }
	next;
    }    
}
my $seq1 = shift;                            #fasta sequence file
my $seq2 = shift;                            #fasta sequence file

unless ($matrixfile and $gap and $outfile)   #complain if options are missing
{ 
    die "Usage: aligner.pl -m <matrixfile> -p <gap penalty (positive)> -a [global | local | noend] -o <outputfile> <seq1 (fasta)> <seq2 (fasta)>\n";
}

my $seq1ref = get_seq($seq1);                #read in sequences
my $seq2ref = get_seq($seq2);

my $matrixref = get_matrix($matrixfile);     #read in score matrix

(my $scoreref, my $traceref, my $maxref) = 
    fill_tables($seq1ref, $seq2ref, $matrixref);  #fill in score and traceback

output_matrix($scoreref, $seq1ref, $seq2ref, "$outfile.score"); 
output_trace($traceref, $seq1ref, $seq2ref, "$outfile.trace");  #print matrices

(my $align1, my $align2) =                   #recursive generation of path
    traceback($traceref, $maxref->[1], $maxref->[2], $seq1ref, $seq2ref);

print "Score: $maxref->[0]\n";
print "\n$align1\n";
print "$align2\n";

#----- fill_tables ----------------------------------------------------#
# for each position in matrix, calculate S_i,j from neighboring cells  #
# plus score for aligning i,j or gap.                                  #
# input: refs to two sequence arrays and matrix of scores for aligning #
# returns: refs to score matrix, trace matrix, and maximum score/pos   #
#----------------------------------------------------------------------#
sub fill_tables
{
    my $seq1 = $_[0];                        #reference to first sequence
    my $seq2 = $_[1];                        #reference to second sequence
    my $matrix = $_[2];                      #reference to score matrix
    my @scoretable;                          #m+1 x n+1; will hold scores
    my @tracetable;                          #m+1 x n+1; will hold path
    my @localmax = (0, 0, 0);                #store score, coordinates of max
    
    if ($global == 1)                        #global alignment
    {
	$scoretable[0][0] = 0;               #set upper left corner to 0
	$tracetable[0][0] = 0;
	for (my $i=1; $i <= ($#$seq1 + 1); $i++)  #index: row; max: seqlen1+1
	{ 
	    $scoretable[$i][0] = -($gap * $i);    #fill in left side
	    $tracetable[$i][0] = '|'; 
	}
	for (my $j=1; $j <= ($#$seq2 + 1); $j++)  #index: col; max: seqlen2+1
	{
	    $scoretable[0][$j] = -($gap * $j);    #fill in the top
	    $tracetable[0][$j] = '-'; 
	}
    }
    else                                     #local; global w/o end gap penalty
    {
	for (my $i=0; $i <= ($#$seq1 + 1); $i++)  #index: row; max: seqlen1+1
	{ 
	    $scoretable[$i][0] = 0;               #fill in left side with 0s
	    $tracetable[$i][0] = 0; 
	}
	for (my $j=0; $j <= ($#$seq2 + 1); $j++)  #index: col; max: seqlen2+1
	{ 
	    $scoretable[0][$j] = 0;               #fill in top row with 0s
	    $tracetable[0][$j] = 0;
	}
    }
    for (my $i=1; $i <= ($#$seq1 + 1); $i++)      #index: row; max: seqlen1+1
    {
	for (my $j=1; $j <= ($#$seq2 + 1); $j++)  #index: col; max: seqlen2+1
	{
	    my $x = $scoretable[$i-1][$j-1] +     #score for aligning i,j
		$matrix->{$seq1->[$i-1]}{$seq2->[$j-1]};
	    
	    my $y = $scoretable[$i-1][$j] - $gap; #score for gap in seq2
	    
	    my $z = $scoretable[$i][$j-1] - $gap; #score for gap in seq1
	    
	    my @tmp;
	    if ($global == 1)                #global - sort by score
	    { @tmp = sort {$b->[0] <=> $a->[0]} ([$x,'\\'],      
						 [$y,'|'],
						 [$z,'-']); }
	    elsif ($global == 2)             #no end gap pen - sort by score
	    { 		
		@tmp = sort {$b->[0] <=> $a->[0]} ([$x,'\\'],
						   [$y,'|'],
						   [$z,'-']); 
		if ( ($i == ($#$seq1 + 1) or $j == ($#$seq2 + 1)) and
		     ($tmp[0][0] >= $localmax[0]) )
		{
		    @localmax = ($tmp[0][0], $i, $j);  #update max score if
		}                                      #in last row/col
	    }                                
	    else                             #local alignment - sort by score
	    { 		
		@tmp = sort {$b->[0] <=> $a->[0]} ([$x,'\\'],
						   [$y,'|'],
						   [$z,'-'],
						   [0,0]); 
		if ($tmp[0][0] >= $localmax[0])
		{
		    @localmax = ($tmp[0][0], $i, $j);  #update overall max
		}
	    }
	    $scoretable[$i][$j] = $tmp[0][0];     #new score = max from choices
	    $tracetable[$i][$j] = $tmp[0][1];     #corresponding direction
	}
    }
    if ($global == 1)
    { 
	@localmax = ( $scoretable[$#$seq1 + 1][$#$seq2 + 1], 
		      $#$seq1 + 1, 
		      $#$seq2 + 1); 
    } 
    return (\@scoretable, \@tracetable, \@localmax);
}

#----- get_matrix ------------------------------------------#
# read in score matrix, in blosum-like format. ignore last  #
# row/column (gap scores). return reference to a hash of    #
# hashes with amino acid or nucleotide letters as keys.     #
#-----------------------------------------------------------#
sub get_matrix
{
    my $file = $_[0];                        # matrix file from commandline
    my %matrix;                              # will be hash of hashes
    my @letters;                             # top row of matrix
    my $len;                                 # index of last thing in letters

    open ( MATRIX, "$file" ) or die "Cannot open score matrix.\n";

    while (<MATRIX>)                         #read each line of score matrix
    {
	chomp;
	if (/^\#/) { next; }                 #skip comments
	if (/^   \w/)                        #row of letters at top
	{ 
	    @letters = /(\w)\s/g;            #read all letters into array
	    $len = $#letters;                
	}
	elsif (/^\w/)                        #row of a letter and scores
	{
	    (my $letter) = /^(\w)/;
	    my @scores = /(-?\d+)/g;         #read scores into array
	    for (my $i=0; $i <= $len; $i++)
	    { $matrix{$letter}{$letters[$i]} = $scores[$i]; }    #fill hashes
	}
    }    
    close MATRIX;
    return \%matrix;
}

#----- get_seq ---------------------------------------------#
# read in sequence from specified sequence file. grab all   #
# characters and put in array; return a ref to the array    #
#-----------------------------------------------------------#
sub get_seq
{
    open ( SEQFILE, $_[0] ) or die "Cannot open sequence file $_[0].\n";
    
    unless (<SEQFILE> =~ /^>/)               #make sure first line starts with '>'
    { die "Sequence file is not in fasta format.\n"; }
    
    my $seq = join('', <SEQFILE>);           #grab the rest of the file
    
    close SEQFILE;

    $seq =~ s/\s//gi;                        #remove whitespace
    my @seq = split('', "\U$seq");           #put sequence in uppercase
                                             #then put it in an array
    return \@seq;
}

#----- output_matrix ---------------------------------------#
# print tab-delimited alignment matrix to file              #
#-----------------------------------------------------------#
sub output_matrix
{
    my $matrixref = $_[0];                   #reference to alignment matrix
    my $s1ref = $_[1];                       #reference to sequence 1
    my $s2ref = $_[2];                       #reference to sequence 2

    open (OUTPUT, ">$_[3]") or die "Can't open output file.\n";

    my $toplabel = join "\t  ", @{$s2ref};   #seq 2 letters across top row
    print OUTPUT "\t\t  $toplabel\n";

    for (my $i=0; $i <= $#$matrixref; $i++)  #go through rows
    {
	if ($i==0) { print OUTPUT "\t"; }
	else { printf OUTPUT "%3s\t", $s1ref->[$i-1]; }     #seq 1 letters 

	for (my $j=0; $j <= $#{$matrixref->[0]}; $j++)      #go through columns
	{
	    printf OUTPUT "%3d\t", $matrixref->[$i][$j];    #alignment scores
	}
	print OUTPUT "\n";
    }	
    close OUTPUT;
}

#----- output_trace ----------------------------------------#
# print traceback matrix to file                            #
#-----------------------------------------------------------#
sub output_trace
{
    my $traceref = $_[0];                    #reference to traceback matrix
    my $s1ref = $_[1];                       #reference to sequence 1
    my $s2ref = $_[2];                       #reference to sequence 2

    open (OUTPUT, ">$_[3]") or die "Can't open output file.\n";

    my $toplabel = join " ", @{$s2ref};      #seq 2 letters across top row
    print OUTPUT "    $toplabel\n";

    for (my $i=0; $i <= $#$traceref; $i++)   #go through rows
    {
	if ($i==0) { print OUTPUT "  "; }
	else { print OUTPUT "$s1ref->[$i-1] "; }  #seq 1 letters down first col
	for (my $j=0; $j <= $#{$traceref->[0]}; $j++)  #go through columns
	{
	    print OUTPUT "$traceref->[$i][$j] ";       #direction
	}
	print OUTPUT "\n";
    }	
    close OUTPUT;
}

#----- traceback -------------------------------------------#
# recursively assemble sequence alignment from trace matrix #
# (slow, but fun). return two sequence strings with gaps.   #
#-----------------------------------------------------------#
sub traceback
{
    my $traces = $_[0];                      #ref to trace matrix
    my $i = $_[1];                           #index of row
    my $j = $_[2];                           #index of column
    my $s1ref = $_[3];                       #ref to sequence 1
    my $s2ref = $_[4];                       #ref to sequence 2

    if ($traces->[$i][$j] eq '0')            #end of alignment
    { 
	$align1 = "";                        #terminate strings
	$align2 = "";
    } 
    elsif ($traces->[$i][$j] eq '\\')        #diagonal: align i, j
    {
	my $a = $s1ref->[$i-1];              #nt/aa at i in seq1
	my $b = $s2ref->[$j-1];              #nt/aa at j in seq2
        #recurse on matrix cell to the left and up
	($align1, $align2) = traceback($traces, $i-1, $j-1, $s1ref, $s2ref);
	$align1 .= $a;
	$align2 .= $b;
    }
    elsif ($traces->[$i][$j] eq '-')
    {
	my $a = "-";                         #put a - to show gap in seq1
	my $b = $s2ref->[$j-1];
        #recurse on matrix cell to the left 
	($align1, $align2) = traceback($traces, $i, $j-1, $s1ref, $s2ref);
	$align1 .= $a;
	$align2 .= $b;
    }
    elsif ($traces->[$i][$j] eq '|')
    {	
	my $a = $s1ref->[$i-1];
	my $b = "-";                         #put a - to show gap in seq2
        #recurse on matrix cell to the top
	($align1, $align2) = traceback($traces, $i-1, $j, $s1ref, $s2ref);
	$align1 .= $a;
	$align2 .= $b;
    }
    return ($align1, $align2);               #return strings of sequences
}
