#! -w
#
#       SAM: Sequence Alignment and Modeling Software System                 
#  (c) 1992-1999 Regents of the University of California, Santa Cruz
#
#          Sequence Alignment and Modeling Software System
#         http://www.cse.ucsc.edu/research/compbio/sam.html
#
# ------------- Citations (SAM, SAM-T99, HMMs) -----------------
# R. Hughey, A. Krogh, Hidden Markov models for sequence analysis:
#  Extension and analysis of the basic method, CABIOS 12:95-107, 1996.
# K. Karplus, C. Barrett, R. Hughey, Hidden Markov models for detecting
#  remote protein homologies, Bioinformatics 14(10):846-856, 1999.
# A. Krogh et al., Hidden Markov models in computational biology:
#  Applications to protein modeling, JMB 235:1501-1531, Feb 1994.
# --------------------------------------------------------------
#
#          Distributed with signed license agreement only.
#  
#          Questions or comments to sam-info@cse.ucsc.edu
#


#!/usr/bin/perl -w

# ncbi-blast-prefilter
#	scores a database (default NR) with ncbi-blastp 2.0
#	and selects out those sequences with E-values less than some
#	threhsold(s) 
#
#
# This script takes one required and several optional arguments:
#  REQUIRED:
#	-out root		creates root_1.fa, root_2.fa, ... for output
#	
#	-seed file.a2m		a guide sequence (or sequences)
#
# OPTIONAL:
#	-prefilter_thresholds 1.e-6,300
#		(maybe try 1.e-3,20,400,3000)
#				specifies the E-value
#				thresholds to use for each output
#				(E-value is expected number of false positives)
#	-db nr			Which non-redundant database to use.
#				Note: database MUST be in FASTA format
#				and have a ncbi-blastp index.
#				E_value thresholds used for each iteration
#				(only first and last used with BLAST)
#	-blast_max_report 10000	Maximum number sequences reported by NCBI-BLAST
#	-keep_temporary	0	don't delete temporary directories
#	-tmp_dir /projects/compbio/tmp	temporary directory


# INSTALLATION DEPENDENCIES:
#	o The first line of the file should give the path name for perl5.
#	o The perl module SamT2k.pm is loaded from the same directory containing
#         this file.
#       o It will in turn load sam-t2k.conf, which has directory
#	  information for programs and files.
#

sub echo_parameters();
sub build_query_files($);
sub blast_query_to_database($$$$);
sub accumulate_blast_hits($);
sub extract_fasta();

{   use FileHandle;
    STDOUT->autoflush(1);
    STDERR->autoflush(1);
    use English;
    use File::Basename;
    use lib dirname($PROGRAM_NAME);

    use SamT2k;

    # Set up defaults
    $keep_temporary=0;
    $which_database = $NR;	# default 
    $blast_max_report = 20000;
    @prefilter_thresholds = (1.0e-6 , 300.);

    &process_command_line;

    $hostname = `hostname`;
    chomp $hostname;
    
    $WORKDIR  = absolute("$tmp_root_dir/blast-prefilter-$hostname-$$");
    make_directories($WORKDIR);
    chdir $WORKDIR;

    echo_parameters();	

    # Should redo build_query files to build one file, do query,
    # accumulate results, and continue
    run_prog("uniqueseq init -alignfile $seed_a2m -a protein -a2mdots 0",
	     "uniqueseq of initial alignment failed");
    build_query_files("init.a2m");
    foreach $query (@queries)
    {   $blast_out = "$query-blast.out";
	blast_query_to_database($ncbi_blast_prog, $query, $which_database, $blast_out);
	accumulate_blast_hits($blast_out);
	if (! $keep_temporary)
	{   run_prog("rm -f $query < /dev/null", "removing blast query file failed");
	    run_prog("rm -f $blast_out < /dev/null", "removing blast output failed");
	}
    }

    print STDERR "# Extracting hits in FASTA format from $which_database\n";
    # HERE WE EXTRACT homolog sequences fron $which_database and put them
    # into files  $root_1.fa, ... $root_n.fa
    
    extract_fasta();
    
    foreach $missing (keys(%blast_hits))
    {   if ($blast_hits{$missing} ne "copied")
	{	print STDERR "Error: $PROGRAM_NAME couldn't find $missing in $which_database\n";
	}
    }
    chdir $tmp_root_dir;
    run_prog("rm -rf $WORKDIR < /dev/null",
    		"removing temporary $WORKDIR failed") if (! $keep_temporary);

    exit 0;
}

#--------------------------------------------------
# subroutines

# parse arguments

sub process_command_line
{   local($i) = 0;

    while ($i <= $#ARGV) 
    {	$_ = $ARGV[$i++];
	if (/^-seed$/)	{   $seed_a2m  = absolute( $ARGV[$i++]);
			    (-e $seed_a2m) ||
				fatal("$seed_a2m doesn't exist");
			}
	elsif (/^-out/)	{   $root = absolute($ARGV[$i++]);	}
	elsif (/^-blast_max_report/)
			{   $blast_max_report =  $ARGV[$i++];	}
	elsif (/^-tmp_dir/)
	{    $tmp_root_dir = absolute($ARGV[$i++]);
	}
	elsif (/^-prefilter_thresholds/)
	{   $prefilter_threshold_spec= $ARGV[$i++];
	    @prefilter_thresholds = split(/,/ , $prefilter_threshold_spec);
	}
	elsif (/^-db/)	{   $which_database = $ARGV[$i++];
			    if ($which_database eq "nr") {$which_database=$NR;}
			    else {$which_database = absolute($which_database);}
			}
	elsif (/^-keep_temporary/)	{   $keep_temporary = $ARGV[$i++]; }
	else	{ fatal("Invalid option '$ARGV[$i-1]'");
		}
    }

    # Make certain that required parameters were provided.
    defined($root) || fatal("Must specify '-out root' with $PROGRAM_NAME");
    if ($#prefilter_thresholds <0)
    {	fatal ("Need at least 1 prefilter_threshold (or omit to get defaults)");
    }
    
    for ($i=0; $i<=$#prefilter_thresholds; $i++)
    {    $out_name[$i] = $root . "_" . ($i+1) . ".fa";
    }

    (-e $which_database) || fatal("Can't find database file $which_database");

    # 10000 is maximum permissible value for E
    for ($i=0; $i<=$#prefilter_thresholds ; $i++)
    {   $blast_e = $prefilter_thresholds[$i];
    	$prefilter_thresholds[$i] = ($blast_e> 10000)? 10000 : $blast_e;
    }
    
    defined($seed_a2m) 
    	|| fatal("Must specify '-seed file.a2m' with $PROGRAM_NAME");
}


sub echo_parameters()
{    # Echo the parameters to stderr
    print STDERR "$PROGRAM_NAME: E-value thresholds used are\n";
    for ($i=0; $i<=$#prefilter_thresholds; $i++)
    {	print STDERR "For $out_name[$i],  $prefilter_thresholds[$i]\n";
    }
    print STDERR "\n";
    print STDERR "$PROGRAM_NAME: db=$which_database\n";
    return;
}


#  accumulate_blast_hits(blast_output)
#	reads the blast_output file and collects ids of hits
#	adding them to hash %blast_hits
#
sub accumulate_blast_hits($)
{   my($blast_out_file) = @_;

    open(BLASTOUT_IN, "<$blast_out_file")
	|| fatal("can't open file $blast_out_file");

    #
    # seek Query line
    #
    my($blastout_in_line);
    my($query_id);
    $blastout_in_line = <BLASTOUT_IN>;
    while ($blastout_in_line
	   and $blastout_in_line !~ /^Query=/)
    {	$blastout_in_line = <BLASTOUT_IN>;
    }
    if ($blastout_in_line
	and $blastout_in_line =~ /^Query=\s+(\S+)/)
    {	$query_id = $1;
    }
    else
    {	print STDERR "Error: file $blast_out_file has no Query line\n";
	return;
    }

    # seek to count of sequences
    $blastout_in_line = <BLASTOUT_IN>;
    while ($blastout_in_line
	   and $blastout_in_line !~ /sequences;/)
    {	$blastout_in_line = <BLASTOUT_IN>;
    }
    $blastout_in_line =~ /^\s*(\S+)/ && ($db_size=$1); # grab first token
    $db_size =~ s/,//;	# remove commas from $db_size
    
    print STDERR "Searched $db_size sequences\n";
    # may want to create a way to return $db_size in a file.
    
    # seek to summary of hits to query
    #
    print STDERR "# parsing ncbi-blastp output\n";
    my($hits_table_hdr);
    $hits_table_hdr = "Sequences producing";
    $blastout_in_line = <BLASTOUT_IN>;
    while ($blastout_in_line
	   and $blastout_in_line !~ /^$hits_table_hdr/)
    {	$blastout_in_line = <BLASTOUT_IN>;
    }

    #
    # skip header and the blank line following it
    #
    my($hits_table_line_pat);
    $hits_table_line_pat = '^(\S+)\s+.*(\d+)\s+(\S+)';
    #			     $1          $2    $3   
    #			     id  comment score E
    $blastout_in_line = <BLASTOUT_IN>;
    $blastout_in_line = <BLASTOUT_IN>;
    while ($blastout_in_line
	   and $blastout_in_line =~ /$hits_table_line_pat/)
    {	my($hitseq_id, $high_score, $evalue);
	$hitseq_id = $1;
	$evalue = $3;
	$evalue = "1.$evalue" if ($evalue =~ /^e/);	# fix up truncated evalues
	$blast_hits{$hitseq_id} = $evalue;
	$blastout_in_line = <BLASTOUT_IN>;
    }
    close BLASTOUT_IN;
}

#
# sub blast_query_to_database($blast_prog,
#			      $query_file, $database_file,
#			      $blast_out_file);
#
sub blast_query_to_database($$$$)
{    my($blast_prog,
       $query_file, $database_file,
       $blast_out_file) = @_;

    # report up to $blast_max_report sequences, but don't report any alignments.
    $blast_params = " -e $prefilter_thresholds[$#prefilter_thresholds]"
    		. " -I -v $blast_max_report -b 0";

    my($blast_cmd) = "$blast_prog -d $database_file -i $query_file"
		     . " $blast_params -o  $blast_out_file";
    run_prog($blast_cmd, "$blast_prog failed");
}

#
# sub build_query_files(database_file)
#	builds query files named 1.fasta, 2.fasta, ...
#	leaves names of files in @queries
#
sub build_query_files($)
{    my($database_file) = @_;

    open(DB_IN, "<$database_file")
	|| fatal("can't open file $database_file");

    my $query_count = 0;
    my($db_in_line);
    my($new_seq_pat);
    $new_seq_pat = '^>(\S+)(.*)';

    $db_in_line = <DB_IN>;
    while ($db_in_line)
    {	my($seq_id, $seq_info, $seq);
	($db_in_line =~ /$new_seq_pat/o)
		|| fatal("bad line in file $database_file");

	$seq_id = $1;
	$seq_info = $2;

	$db_in_line = <DB_IN>;
	$seq = "";
	while ($db_in_line
	       and $db_in_line !~ /$new_seq_pat/o)
	{   $seq .= $db_in_line;
	    chomp($seq);
	    $db_in_line = <DB_IN>;
	}

	$query_count ++;
	my($query_out_file) = $query_count . ".fasta";
	push @queries, $query_out_file;
	open(QUERY_OUT, ">$query_out_file")
	    || fatal("can't open file $query_out_file");
	print QUERY_OUT ">$seq_id\n";
	print QUERY_OUT "$seq\n";
	close (QUERY_OUT);
    }
    close (DB_IN);
}

sub extract_fasta()
{
    #BUG: this should work ok if the database has NCBI sequence naming conventions
    # 	but if the database has its own names, then fastacmd will do
    #	name mangling, which will ruin the extraction.
    
    # use fastacmd to get the largest of the sets of sequences
    open(ALL_IDS, "| sort >all.ids");
    foreach $id (keys(%blast_hits))
    {    print ALL_IDS "$id\n";
    }
    close ALL_IDS;
    
    $big_file = $out_name[$#prefilter_thresholds];
    run_prog("$ncbi_fastacmd_prog -d $which_database -i all.ids > $big_file");
    
    for ($i=0; $i< $#prefilter_thresholds; $i++)
    {   $file_ref[$i]  = new IO::File;
	open($file_ref[$i], ">$out_name[$i]")
	    || fatal("Can't open $out_name[$i]: $!");
	$copying[$i] =0;
    }
    open(BIG_IN, "< $big_file") || fatal("Can't read $big_file");

    while(<BIG_IN>)
    {   if (/^>\s*(\S+)[, ]?/)
	{   # new id
	    $id = $1;
	    $hit_found = defined $blast_hits{$id};
	    if (!$hit_found)
	    {   $id =~ /^(gi\|\d+)|/;
		$short_id = $1;
		if (defined $short_id && defined $blast_hits{$short_id})
		{	 # BLAST truncated the identifier
		    $id=$short_id;
		    $hit_found=1;
		}
	    }
	    
	    for ($i=0; $i< $#prefilter_thresholds; $i++)
	    {    $copying[$i] = $hit_found &&
		    ($blast_hits{$id} <= $prefilter_thresholds[$i]);
	    }
	    $copying[$#prefilter_thresholds] = $hit_found;
	     
	    if ($hit_found)
	    {   $blast_hits{$id} = "copied";
	    }
	}
	for ($i=0; $i< $#prefilter_thresholds; $i++)
	{    print {$file_ref[$i]}    $_ if $copying[$i];
	}
    }
    for ($i=0; $i< $#prefilter_thresholds; $i++)
    {    close $file_ref[$i];
    }
    
}

sub old_extract_fasta()
{
    # BUG: FOR NCBI-BLAST it would be much more efficient to use 
    #	fastacmd -i
    for ($i=0; $i<= $#prefilter_thresholds; $i++)
    {   $file_ref[$i]  = new IO::File;
	open($file_ref[$i], ">$out_name[$i]")
	    || fatal("Can't open $out_name[$i]: $!");
	$copying[$i] =0;
    }
    open(NRP_IN, "< $which_database") || fatal("Can't read $which_database");

    while(<NRP_IN>)
    {   if (/^>\s*(\S+)[, ]?/)
	{   # new id
	    $id = $1;
	    $hit_found = defined $blast_hits{$id};
	    if (!$hit_found)
	    {   $id =~ /^(gi\|\d+)|/;
		$short_id = $1;
		if (defined $short_id && defined $blast_hits{$short_id})
		{	 # BLAST truncated the identifier
		    $id=$short_id;
		    $hit_found=1;
		}
	    }
	    
	    for ($i=0; $i< $#prefilter_thresholds; $i++)
	    {    $copying[$i] = $hit_found &&
		    ($blast_hits{$id} <= $prefilter_thresholds[$i]);
	    }
	    $copying[$#prefilter_thresholds] = $hit_found;
	     
	    if ($hit_found)
	    {   $blast_hits{$id} = "copied";
	    }
	}
	for ($i=0; $i<= $#prefilter_thresholds; $i++)
	{    print {$file_ref[$i]}    $_ if $copying[$i];
	}
    }
    for ($i=0; $i<= $#prefilter_thresholds; $i++)
    {    close $file_ref[$i];
    }
}
