#!/lab/bin/perl

# $header$

# motifBS.pl - motif Builder and Searcher
# (c) 2003 Ed Green, UC Regents

use Inline C;
use Getopt::Std;
use Bio::AlignIO;
use Bio::SeqIO;
use strict;
use vars qw( $opt_a $opt_o $opt_d $opt_m $opt_c $opt_s $opt_b $opt_z $opt_B $opt_e );

my %base_code = ( 'A' => 0,
		  'C' => 1,
		  'G' => 2,
		  'T' => 3,
		  'U' => 3,
		  'a' => 0,
		  'c' => 1,
		  'g' => 2,
		  't' => 3,
		  'u' => 3,
		  '-' => 4 );

my @pc = ( 1, 1, 1, 1 ); # pseudo-counts

MAIN : {
    my ( $log_table_p, $aln, $freq_p, $bkgd_p, $search_db );

    &init();                       # get options or show help

    # User-supplied PSSM
    if ( $opt_o ) {
	$log_table_p = &parseMat();
    }
    
    # Make PSSM from MSA
    else {
	$aln = &getAlign();         # open fasta MSA file and return
    	                            # bioperl align object

	$freq_p = &getFreq( $aln ); # go through the sequences in the
	                            # MSA and make an array of arrays
                                    # with frequency information 

	$bkgd_p = &getBkgd( $freq_p);  # go through the file of sequences
                                       # to get the background frequencies
    
	$log_table_p = &makeLogTable( $freq_p, $bkgd_p );

	&outputPSSM( $log_table_p );     # make meme-style output using the
                                         # log-odds table
    }    

    # Search Database
    if ( $opt_d ) {                   # if the user specified a search database
	$search_db = &getSearchDb();  # then search it with the PSSM
	&searchDb( $search_db, $log_table_p );
    }
}

# Parse meme-style PSSM data into log-odds PSSM table
sub parseMat {
    open( MAT, $opt_o ) or die( "$opt_o: $!\n" );
    my ( @log_table, @el );
    my ( $l, $j );
    my $i = 0;

    <MAT>;
    <MAT>;

    while( $l = <MAT> ) {
	$i++;
	$l =~ s/^\s+//;
	@el = split( /\s+/, $l );
	for( $j = 0; $j <= 3; $j++ ) {
	    $log_table[ $i ][ $j ] = $el[ $j ];
	}
    }
    close( MAT );
    return \@log_table;
}

# Args    => Bio::SeqIO object (search database),
#            pointer to log-table PSSM
# Returns => nothing 
# Aligns PSSM to every possible position in the database of sequences.
# Generates a list of hits above some cutoff score.  Sends this list
# to &output() to be sent to STDOUT.
sub searchDb {
    my $sdb = shift;  # search database (Bio::SeqIO)
    my $lt_p = shift; # log table of PSSM
    my $motif_len = $#{ $lt_p }; # zero-indexed length of PSSM
    my $hits_p;
    my ( $seq, $length, $seq_str, $seq_id, $hit_p, $hits_p );
    my ( $length, $seq_str, $seq_id, $score, $i, $start, $score );
    my ( @unraveled );

    # If Big database option, go ahead and print the header and then
    # we'll just show the hits as we see them
    &outputHeader( $lt_p, $motif_len );

    if ( $opt_z == 2 ) {
	@unraveled = &unravelMotif( $lt_p, $motif_len );
    }

    while ( $seq = $sdb->next_seq() ) {
	$length = $seq->length();      # length of sequence
	next if $length == 0;
	$seq_str = $seq->seq();        # sequence string
	$seq_id  = $seq->display_id(); # sequence identifier

	if ( $opt_e ) {
	    $hits_p = &setUpHits();
	}

	if ( $length < $motif_len ) {
	    for ( $i = 1; $i <= ($motif_len - $length + 1); $i++ ) {
		$score = &getScoreSeq( $seq_str, 1, $length,
				       $lt_p, $i, $length );
		if ( $score >= $opt_c ) {
		    $hit_p = [ $seq_id,
			       $length,
			       &trunc( $score ), 
			       &trunc( $score / &logN( $length, $opt_s ) ),
			       $seq_str,
			       1, $length ];
		    $hits_p = &addHit( $hits_p, $hit_p );
		}
	    }
	}

	elsif ( $length == $motif_len ) {
	    $score = &getScoreSeq( $seq_str, 1, $length,
				   $lt_p, 1, $length );
	    if ( $score >= $opt_c ) {
		$hit_p = [ $seq_id, $length, &trunc( $score ),
			   &trunc( $score / &logN( $length, $opt_s ) ),
			   $seq_str, 1, $length ];
		$hits_p = &addHit( $hits_p, $hit_p );
	    }
	}

	else {
	    if ( $opt_z != 2 ) {
		for ( $i = 1; $i <= ($length - $motif_len + 1); $i++ ) {
		    $score = &getScoreSeq( $seq_str, $i, $motif_len,
					   $lt_p, 1, $motif_len );
		    if ( $score >= $opt_c ) {
			$hit_p = [ $seq_id, $length, &trunc( $score ),
				   &trunc( $score / &logN( $length, $opt_s ) ),
				   substr( $seq_str, ($i - 1), $motif_len ), # the sequence hit
				   $i, ( $i + $motif_len - 1 ) ];
			$hits_p = &addHit( $hits_p, $hit_p );
		    }
		}
	    }
	    else { # C implementation
		my @hs = &search( $seq_str, $opt_c, $motif_len, @unraveled );
		for ( my $i = 0; $i <= $#hs; $i += 2 ) {
		    $start = $hs[ $i ] + 1;
		    $score = $hs[ $i + 1 ];
		    $hit_p = [ $seq_id, $length, &trunc( $score ),
			       &trunc( $score / &logN( $length, $opt_s ) ),
			       substr( $seq_str, $start - 1, $motif_len ),
			       $start, ( $start + $motif_len - 1 ) ];
		    $hits_p = &addHit( $hits_p, $hit_p );
		}
	    }
	}
    
	# now %hits is filled up with all the hits on this sequence that
	# were better than the cutoff, $opt_c
	next if $opt_B;	
	next if !$opt_e;
	foreach $hit_p ( @$hits_p ) {
	    &outputHit( $hit_p );
	}
    }

    unless ( $opt_B || $opt_e ) {
	&outputHitsInOrder( $motif_len, $lt_p, $hits_p );
    }
}

sub setUpHits {
    my $hits_p = [ ['dummy', 'dummy', $opt_c, 'dummy', 
		    'dummy', 'dummy', 'dummy'] ];
    return $hits_p;
}

sub unravelMotif {
    my $ltp   = shift; # log-table pointer
    my $m_len = shift; # motif-length
    my @unraveled;
    
    for( my $i = 1; $i <= $m_len; $i++ ) {
	for ( my $j = 0; $j <= 3; $j++ ) {
	    push( @unraveled, $ltp->[$i][$j] );
	}
    }
    return @unraveled;
}

sub logN {
    my $num = shift;
    my $base = shift;
    if ( $num <= 0 ) {
	warn( "Cannot take log of $num\n" );
	return;
    }
    return log($num)/log($base);
}

sub outputHitsInOrder {
    my $motif_len   = shift;
    my $log_table_p = shift;
    my $hits_p      = shift;
    
    my ( $score, $hit );

    foreach $score ( sort { $b <=> $a } keys %{ $hits_p } ) {
	foreach $hit ( @{ $hits_p->{ $score } } ) {
	    &outputHit( $hit );
	}
    }
}

sub outputHeader {
    my $log_table_p = shift;
    my $motif_len   = shift;
    print( "# PSSM source: ", &PSSMSource(), "\n" );
    print( "# PSSM length: $motif_len\n" );
    print( "# Max score: ", &maxScore( $log_table_p ), "\n" );
    print( "# Min score: ", &minScore( $log_table_p ), "\n" );
    print( "# Avg score: ", &avgScore( $log_table_p ), "\n" );
    print( "# Log base: $opt_s\n" );
    print( "# Background base frequencies: ", &showBkgd(), "\n" );
    print( "\#\#\# Hits \#\#\#\n" );
    print( "### Sequence-id length raw-score length-normalized-score hit-sequence start end\n" );
}

sub addHit {
    my $hits_p = shift;
    my $hit_p  = shift;
    my ( $ins_pos );
    my ( @new_hits );
    if ( $opt_B ) {
	&outputHit( $hit_p );
	return $hits_p;
    }
    
    # User wants only top e hits
    # data structure is different (a sorted array) in this case
    if ( $opt_e ) {
	# This hit is in the top e hits thusfar
	if ( $hit_p->[ 2 ] > $hits_p->[ $#{ $hits_p } ]->[ 2 ] ) {
	    $ins_pos = &findInsPos( $hit_p->[ 2 ], $hits_p );
	    @new_hits = ( @{ $hits_p }[ 0..($ins_pos-1)], 
			  $hit_p,
			  @{ $hits_p }[$ins_pos..(&min($opt_e-2, $#{$hits_p}))] );
	    $hits_p = \@new_hits;
	}
	return $hits_p;
    }

    # gotta remember everything
    else { 
	push( @{ $hits_p->{ $hit_p->[ 2 ] } }, $hit_p );
	return $hits_p;
    }
}

sub findInsPos {
    my $hit_score = shift;
    my $hits_p    = shift;
    my $pos       = 0;
    my ( $hit_p );
    foreach $hit_p ( @{$hits_p} ) {
	if ( $hit_score > $hit_p->[ 2 ] ) {
	    return $pos;
	}
	else {
	    $pos++;
	}
    }
    return $pos;
}

sub outputHit {
    my $hit_p = shift;
    unless( $hit_p->[ 0 ] eq 'dummy' ) {
	print( join( ' ', @{ $hit_p } ), "\n" );
    }
}

# Finds and returns the source of this PSSM - either a pre-existing 
# PSSM or a MSA file
sub PSSMSource {
    if ( $opt_o ) {
	return $opt_o;
    }
    else {
	return $opt_a;
    }
}

sub avgScore {
    my $ltp = shift; # log-table pointer
    my ( $i, $j, $pos_total, $total );
    for ( $i = 1; $i <= $#{ $ltp }; $i++ ) {
	$pos_total = 0;
	for ( $j = 0; $j <= 3; $j++ ) {
	    $pos_total += $ltp->[ $i ][ $j ];
	}
	$total += $pos_total / 4;
    }
    return $total / $#{ $ltp };
}

sub maxScore {
    my $ltp = shift; # log-table pointer
    my ( $i, $j, $max );
    for ( $i = 1; $i <= $#{ $ltp }; $i++ ) {
	my $max_at_this_position = 0;
	for ( $j = 0; $j <= 3; $j++ ) {
	    if ( $ltp->[ $i ][ $j ] > $max_at_this_position ) {
		$max_at_this_position = $ltp->[ $i ][ $j ];
	    }
	}
	$max += $max_at_this_position;
    }
    return $max;
}

sub minScore {
    my $ltp = shift; # log-table pointer
    my ( $i, $j, $min );
    for ( $i = 1; $i <= $#{ $ltp }; $i++ ) {
	my $min_at_this_position = 0;
	for ( $j = 0; $j <= 3; $j++ ) {
	    if ( $ltp->[ $i ][ $j ] < $min_at_this_position ) {
		$min_at_this_position = $ltp->[ $i ][ $j ];
	    }
	}
	$min += $min_at_this_position;
    }
    return $min;
}
	    
sub showBkgd {
    if ( $opt_b ) {
	return $opt_b;
    }
    else {
	return ( "flat" );
    }
}

# getScoreSeq
# Returns the single score of the input sequence against the PSSM by
# aligning the PSSM for the defined start position to the defined start
# position on the sequence
sub getScoreSeq {
    my $seq    = shift; # sequence being searched
    my $seqs   = shift; # start position on sequence being searched
    my $seql   = shift; # length of region to score
    my $pssm_p = shift; # pointer to PSSM
    my $pssms  = shift; # start position on PSSM
    my $pssml  = shift; # length of PSSM region to score

    my $subseq = substr( $seq, $seqs - 1, $seql );
    
    my ( $score_region_length, $cur_char, $i );
    my $score = 0;
    
    if ( length( $subseq ) < $pssml ) {
	# sequence is shorter than pssm
	$score_region_length = length( $subseq );
    }
    else {
	$score_region_length = $pssml;
    }

    for ( $i = 1; $i <= $score_region_length; $i++ ) {
	$cur_char = substr( $subseq, ($i - 1), 1 );
	$score += &getScoreChar( $pssm_p, $i, $pssms, $cur_char ); 
    }
    return $score;
}

sub getScoreChar {
    my $pssm_p   = shift;
    my $i        = shift;
    my $pssms    = shift;
    my $cur_char = shift;
    if ( $cur_char eq 'A' || $cur_char eq 'C' ||
	 $cur_char eq 'G' || $cur_char eq 'T' ||
	 $cur_char eq 'U' ) {
	return $pssm_p->[ $pssms + $i - 1 ][ $base_code{ $cur_char } ];
    }
    else {
	return 0;
    }
}

sub getSearchDb {
    my $search_db = Bio::SeqIO->new( -file => $opt_d,
				     '-format' => 'Fasta' );
    return $search_db;
}

sub getAlign {
    my $aln_io = Bio::AlignIO->new( -file => $opt_a,
				    '-format' => 'fasta' );
    my $aln = $aln_io->next_aln();
    return $aln;
}

# Arg     => Bio::Align object
# Returns => frequency table of bases
sub getFreq {
    my $aln = shift;
    my @freq;

    my ( $seq, $length, $i, $cur );

    foreach $seq ( $aln->each_seq ) {
	$length = $seq->length;
	for ( $i = 1; $i <= $length; $i++ ) {
	    $cur = $seq->subseq( $i, $i );
	    if ( $cur eq 'A' || $cur eq 'C' || $cur eq 'G' ||
		 $cur eq 'T' || $cur eq 'U' || $cur eq 'a' ||
		 $cur eq 'c' || $cur eq 'g' || $cur eq 't' || 
		 $cur eq 'u' ) {
		$freq[ $i ][ $base_code{ $cur } ]++;
		# Find the symbol in the ith position in this sequence.
		# Then, using the %base_code, increment the appropriate
		# array element of the ith array of @freq
	    }
	}
    }
    return \@freq;
}

sub getBkgd {
    my $freq_p = shift;
    my @bkgd_freq;
    my ( $length, $i, $j, $cur );
    my $obs_len = 0; # the observed length of all these sequences

    # Find background frequencies from a file of sequences?
    if ( $opt_b ) {
	# Same file as was used to make position-specific counts?
	if ( $opt_b eq $opt_a ) {
	    for ( $i = 1; $i <= $#{ $freq_p }; $i++ ) {
		for ( $j = 0; $j <= 3; $j++ ) {
		    $obs_len += $freq_p->[ $i ][ $j ];
		    $bkgd_freq[ $j ] += $freq_p->[ $i ][ $j ];
		}
	    }
	}
	# different file
	else {
	    my $seqio = Bio::SeqIO->new( -file => $opt_b,
					 '-format' => 'Fasta' );
	    while ( my $seq = $seqio->next_seq() ) {
		$length = $seq->length();
		for ( $i = 1; $i <= $length; $i++ ) {
		    $cur = $seq->subseq( $i, $i );
		    if ( $cur eq 'A' || $cur eq 'C' || $cur eq 'G' ||
			 $cur eq 'T' || $cur eq 'U' || $cur eq 'a' ||
			 $cur eq 'c' || $cur eq 'g' || $cur eq 't' || 
			 $cur eq 'u' ) {
			$bkgd_freq[ $base_code{ $cur } ]++;
			$obs_len++;
		    }
		}
	    }
	}

	return ( [ ($bkgd_freq[0]/$obs_len),
		   ($bkgd_freq[1]/$obs_len),
		   ($bkgd_freq[2]/$obs_len),
		   ($bkgd_freq[3]/$obs_len) ] );
    }

    # Flat background
    else {
	return ( [ 0.25, 0.25, 0.25, 0.25 ] );
    }
}

sub makeLogTable {
    my $freq_p = shift;
    my $bkgd_p = shift;
    my @log_table;
    my ( $total_count_i, $i, $j );

    my $total_pc = $pc[ 0 ] + $pc[ 1 ] + $pc[ 2 ] + $pc[ 3 ];

    for ( $i = 1; $i <= $#{ $freq_p }; $i++ ) {
	# Find the total count at this position in the MSA (may be
	# less than total number of sequences because of gaps
	$total_count_i = ( $freq_p->[$i][0] + $freq_p->[$i][1]
			   + $freq_p->[$i][2] + $freq_p->[$i][3] );
	for ( $j = 0; $j <= 3; $j++ ) {
	    $log_table[ $i ][ $j ] 
		= &logN( (($freq_p->[ $i ][ $j ] + $pc[ $j ]) / 
			  ($total_count_i + $total_pc)) /
			 $bkgd_p->[ $j ], $opt_s );
	}
    }
    return \@log_table;
}

sub outputPSSM {
    if ( $opt_m ) {
	my $log_table_p = shift;
	my $i;
	open ( PSSM, ">$opt_m" ) or die( "$opt_m: $!\n" );
	print PSSM ( "ALPHABET= ACGT\n" );
	print PSSM ( "log-odds matrix: alength= 4 w= " );
	print PSSM ( $#{ $log_table_p }, "\n" );
	for ( $i = 1; $i <= $#{ $log_table_p }; $i++ ) {
	    print PSSM ( " " );
	    print PSSM ( &trunc( $log_table_p->[ $i ][ 0 ], "  " ) );
	    print PSSM ( &trunc( $log_table_p->[ $i ][ 1 ], "  " ) );
	    print PSSM ( &trunc( $log_table_p->[ $i ][ 2 ], "  " ) );
	    print PSSM ( &trunc( $log_table_p->[ $i ][ 3 ], "\n" ) );
	}
	close ( PSSM );
    }
}
     

sub trunc {
    my $n = shift;
    my $last = shift;
    
    my ( $before_dot, $after_dot );
   
    if ( $n =~ /\./ ) {
	( $before_dot, $after_dot ) = ( $n =~ /(-*\d+)\.(\d+)/ ); 
	$after_dot = substr( $after_dot, 0, 4 );
	return ( "$before_dot.$after_dot$last" );
    }
    else {
	return ( "$n$last" );
    }
}

sub min {
    my $min = $_[ 0 ];
    foreach my $val ( @_ ) {
	if ( $val < $min ) {
	    $min = $val;
	}
    }
    return $min;
}

sub init {
    getopts( 'a:d:m:c:s:b:o:z:B:e:' );
    unless ( $opt_a || $opt_o ) {
	print( "motifBS.pl - Builds and searches with a DNA seqeunce motif.\n" );
	print( "Converts a multiple sequence alignment to a PSSM\n" );
	print( "in MAST output format.  Then, can search a database\n" );
	print( "of sequences with the PSSM.\n" );
	print( "The alphabet for motifs is 'A', 'C', 'G', 'T', and 'U'\n" );
	print( "Other characters, like '-' and 'N', are treated as no-data\n" );
	print( "in constructing the motifs and given score of 0 in\n" );
	print( "searching a database.\n" );
	print( "OPTIONS:\n" );
	print( "-a [ MSA in fasta format ]\n" );
	print( "-o [ Meme log-odds matrix (previously computed) ]\n" );
	print( "-d [ Fasta database to search with the motif (opt.) ]\n" );
	print( "-m [ Filename to write out the log-odds matrix (opt.) ]\n" );
	print( "-c [ Score cutoff for reporting scores; default = 10 ]\n" );
	print( "-s [ Base units for log-odds score matrix; default = 2 ]\n" );
	print( "-b [ Fasta database of sequences from which to\n" );
	print( "     calculate the background frequencies.  If\n" );
	print( "     none is specified, then\n" );
	print( "     f(A) = f(T) = f(G) = f(C) = 0.25 is used ]\n" );
	print( "-z [ 1 => perl search implementation (default);\n" );
	print( "     2 => c search implementation ]\n" );
	print( "-B [ Big database option.  Runs fast when searching large\n" );
	print( "     databases (>50,000 sequences), but gives hits in order\n" );
	print( "     they are found, not in order of decreasing score.\n" );
	print( "     Default = 0 (off) ]\n" );
	print( "-e [ If there are multiple hits of the profile against a\n" );
	print( "     sequence above the score cutoff, just show this many\n" );
	print( "     of the best ones. ]\n" );
	print( "Note that -a and -o are mutually exclusive.  If both are\n" );
	print( "specified, then -o takes precedence.\n" );
	print( "Also, -B and -e are mutually exclusive. If both are specified\n" );
	print( "then -B takes precedence.\n" );
	exit( 0 );
    }
    
    unless( $opt_c ) {
	$opt_c = 10;
    }
    unless ( $opt_s ) {
	$opt_s = 2;
    }
    unless ( $opt_z ) {
	$opt_z = 1;
    }
    unless ( $opt_B ) {
	$opt_B = 0;
    }
    unless ( $opt_e ) {
	$opt_e = 0;
    }
}


#ALPHABET= ACGT
#log-odds matrix: alength= 4 w= 9
# -4.275  -0.182  -4.195   1.408
# -4.296  -1.487   1.880  -0.816
# -2.160  -1.492  -4.171   1.474
# -0.810  -4.076   1.872  -2.164
# 1.537  -1.487  -4.195  -4.205
# 0.113   0.340  -0.237  -0.209
# -0.454   0.923   0.390  -0.834
# -1.336  -0.082   0.905   0.100
# 0.674  -4.183   0.130  -0.201
# log-odds matrix: alength= 4 w= 6
# -2.032   0.324   1.371  -0.781
# -0.409   0.560  -0.250   0.119
# -4.274  -0.519  -0.260   1.167
# -2.188   2.300  -4.191  -2.465
# 1.265  -4.111  -0.267  -2.180
# -1.977   2.158  -1.661  -2.071 


__END__
__C__
double mot_score( double motif[100][4], int mot_len, char* seq ) {
    double tot = 0.0;
    int i;
    for( i = 0; i < mot_len; i++ ) {
	switch ( seq[i] ) {
	    case 'A': {
		tot += motif[i][0];
		break;
	    }
	    case 'a': {
		tot += motif[i][0];
		break;
	    }
	    case 'C': {
		tot += motif[i][1];
		break;
	    }
	    case 'c': {
		tot += motif[i][1];
		break;
	    }
	    case 'G': {
		tot += motif[i][2];
		break;
	    }
	    case 'g': {
		tot += motif[i][2];
		break;
	    }
	    case 'T': {
		tot += motif[i][3];
		break;
	    }
	    case 't': {
		tot += motif[i][3];
		break;
	    }
	    case 'U': {
		tot += motif[i][3];
		break;
	    }
	    case 'u': {
		tot += motif[i][3];
		break;
	    }
	}
    }
    return tot;
}

void search( char* seq, double cutoff, int mot_len, SV* mot, ... ) {
    Inline_Stack_Vars;
    double motif[100][4];
    int i, j, seq_len;
    double cur_score;

    if ( mot_len > 100 ) {
	printf( "Motif is too long - can't be longer than 100 positions\n" );
	return 1;
    }

    for ( i = 0; i < mot_len; i++ ) {
	for ( j = 0; j <= 3; j++ ) {
	    motif[i][j] = SvNV(Inline_Stack_Item((i*4)+j+3));
	}
    }

    seq_len = strlen( seq );

    Inline_Stack_Reset;

    for ( i = 0; i < seq_len - mot_len + 1; i++ ) {
       	cur_score = mot_score( motif, mot_len, &seq[i] );
	if ( cur_score >= cutoff ) {
	    Inline_Stack_Push(sv_2mortal(newSViv(i)));
	    Inline_Stack_Push(sv_2mortal(newSVnv(cur_score)));
	    Inline_Stack_Done;
	}
    }
    Inline_Stack_Done;
}
