#!/usr/local/bin/perl5.6.0 -w

BEGIN {
  unshift(@INC, "/nfs/disk222/yeastpub/zoo/modules")
};

use strict;
use Getopt::Std;
use File::Basename;
use POSIX;
use AnT::Searchfactory;

# $Header: /nfs/disk222/yeastpub/Repository/zoo/general/big_blast.pl,v 1.23 2000/11/27 09:44:12 kmr Exp $

# purpose: This script will chop up a large sequence, run blast on each bit and
# purpose: then write out an EMBL feature table and a MSPcrunch -d file
# purpose: containing the hits.

my $piece_size = 50000;
my $min_score = 0;

sub usage
{
  die <<EOF;
usage: $0 [-l] [-t|-x|-tx] [-S size] [-j num] [-M num] database input_file [blast options]
   or: $0 [-M num] blast_output_file
   or: $0 -h (for help)

  -l      use lsrun to share the blast jobs over multiple machines
  -x      selects blastx instead of blastn
  -t      selects tblastn instead of blastn
  -tx     selects tblastx instead of blastn
  -j num  sets the number of blast jobs to run simultaneously
  -M num  sets the minimum score for the output hits (default $min_score)
  -S num  sets the maximum size of sequences sent to blast (default $piece_size)
  -w      use wublast instead of blast1.4
  -2      use blast2 not blast1.4

The first form of the command will run blast on the given input_file
against the given database.  The -S flag sets the size of the pieces of
sequence that are passed to blast.  All arguments after the input_file
parameter will be passed directly to blast.

The second form reads and processes the given blast_output_file to
produce an MSPcrunch crunch file and an EMBL feature table.

In both cases there are two output files.  One is an EMBL feature
table with one feature per hit.  The other is a file in the
format generated by MSPcrunch -d.
EOF
}

my $blast2_program = "blastall_2.0.11";
my $blast_program = "blastn";

my $lsrun = "lsrun";
my $lsload = "lsload";
my $child_count = 1;

if (!@ARGV) {
  warn "error: not enough arguments\n\n";
  usage;
}

if ($ARGV[0] eq "-h") {
  usage;
}

my $version_flag = "";

my $use_blast2 = 0;
my $use_wublast = 0;
my $use_lsrun = 0;

# if true some jobs (50%) will be sent to plato
my $use_plato = 0;

my %options = ();

getopts ("lpw2xtj:M:S:", \%options);

if (exists $options{2} && exists $options{w}) {
  die "$0: error the -2 and -w flags are mutually exclusive\n";
}

if (exists $options{j}) {
  if (@ARGV < 2) {
    warn "error: not enough arguments\n\n";
    usage;
  }

  $child_count = $options{j};
}

if (exists $options{M}) {
  $min_score = $options{M};
  $version_flag .= "M${min_score}_";
}

if (exists $options{S}) {
  if (@ARGV < 2) {
    warn "error: not enough arguments\n\n";
    usage;
  }

  $piece_size = $options{S};
}

if (exists $options{2}) {
  if (@ARGV < 2) {
    warn "error: not enough arguments\n\n";
    usage;
  }

  $use_blast2 = 1;
}

if (exists $options{l}) {
  if (@ARGV < 2) {
    warn "error: not enough arguments\n\n";
    usage;
  }

  $use_lsrun = 1;
}

if (exists $options{'x'}) {
  if (@ARGV < 2) {
    warn "error: not enough arguments\n\n";
    usage;
  }

  $blast_program =~ s/n$/x/;
  $version_flag .= "x_";
}

if (exists $options{t}) {
  if (@ARGV < 2) {
    warn "error: not enough arguments\n\n";
    usage;
  }

  $blast_program = 't' . $blast_program;
  $version_flag .= "t_";
}

if (exists $options{w}) {
  if (@ARGV < 2) {
    warn "error: not enough arguments\n\n";
    usage;
  }

  $use_wublast = 1;
  $blast_program = 'wu' . $blast_program;
}

if (exists $options{p}) {
  if (@ARGV < 2) {
    warn "error: not enough arguments\n\n";
    usage;
  }

  $use_plato = 1;
}


if (@ARGV == 1) {
  unlink "$ARGV[0].crunch", "$ARGV[0].tab";

  process_blast ($ARGV[0], 0);

  print STDERR "wrote $ARGV[0].crunch and $ARGV[0].tab\n";
  exit;
}

if (@ARGV < 1) {
  warn "error: not enough arguments\n\n";
  usage;
}

my @blast_options;
#my @blast_options = qw { -span1 V=2000 B=2000 };

if ($use_blast2) {
  @blast_options = ("-p", $blast_program);
  $version_flag .= "2";
} else {
  if ($use_wublast) {
    @blast_options = qw { };
    $version_flag .= "w";
  } else {
    @blast_options = qw { B=1000000 -span1 V=1000000 -H=0 -hspmax 999999 };
    $version_flag .= "1";
  }
}


my $database = shift;
my $in_file = shift;

if (!open (IN, "$in_file")) {
  if (!open (IN, "$in_file.seq")) {
    die "can't open $in_file or $in_file.seq\n";
  } else {
    print STDERR "note: opened $in_file.seq\n";
  }
}

my $database_name = basename($database);

my $base_name = "big_blast_" . $version_flag . "_" . $piece_size .
                "_${database_name}_$in_file";

push @blast_options, @ARGV;

# the files that contain the small pieces of sequence
my @piece_files = ();
my @piece_offsets = ();

sub make_piece
{
  my ($sequence, $offset) = ($_[0], $_[1] + 1);
  my $out_file = sprintf ("$base_name.%8.8d", $offset);

  $sequence =~ s/(.{60})/$1\n/g;

  if (-s $out_file) {
    print STDERR "not writing $out_file - file exists\n";
    return $out_file;
  }

  open OUT, ">$out_file" or die "$0: can't open $out_file\n";

  print OUT ">$out_file\n";
  print OUT $sequence,"\n";

  close OUT;

  $out_file;
}

my $sequence = "";

while (<IN>) {
  if (!/^>/) {
    chomp;
    $sequence .= $_;
  }
}

my $bases_so_far = 0;

while ($bases_so_far < length $sequence) {
  my $piece = substr $sequence, $bases_so_far, $piece_size;
  push @piece_files, (make_piece $piece, $bases_so_far);
  push @piece_offsets, $bases_so_far;

  $bases_so_far += length $piece;
}

print STDERR "sequence split into ", scalar (@piece_files), " pieces\n";

my @children = ();

# returns true iff the given file seems like a valid blast output file
sub check_blast_output
{
  my ($in_file) = @_;

  if (!open IN_FILE, $in_file) {
    return 0;
  }

  local $/ = undef;

  my $file_size = -s IN_FILE;

  if (!$file_size) {
    close IN_FILE;
    print STDERR "$in_file exists but is truncated\n";
    return 0;
  }

  my $tail_size = 10000;

  if ($file_size > $tail_size) {
    if (!seek IN_FILE, $file_size - $tail_size, SEEK_SET) {
      warn "can't seek on $in_file\n";
      close IN_FILE;
      return 0;
    }
  }

  my $end_of_file = <IN_FILE>;

  if ($end_of_file =~ m/effective length of database|Total cpu time/) {
    close IN_FILE;
    return 1;
  } else {
    close IN_FILE;
    print STDERR "$in_file exists but is truncated\n";
    return 0;
  }
}

sub process_blast
{
  my ($file, $offset) = @_;

  my $crunch_file = "$file.crunch";
  my $tab_file = "$file.tab";

  if (-s $crunch_file && -s $tab_file) {
    print STDERR "skipping $crunch_file and $tab_file - file exists\n";
    return;
  }

  open CRUNCH_OUT, ">$crunch_file" or die "couldn't open $crunch_file\n";
  open TAB_OUT, ">$tab_file" or die "couldn't open $tab_file\n";

  print STDERR "opened $crunch_file for writing\n";

  my $blast_obj;

  eval {
#      $blast_obj = Bio::Tools::Blast->new (-file   => "$file",
#                                           -parse  => 1);
    $blast_obj = AnT::Searchfactory->make (-file   => "$file",
                                           -program  => "blast");
  };

  if ($@) {
    print STDERR $@;
    return;
  }

  my $first_results = $blast_obj->next_result;

#  $first_results->store_align (0);

  if (!defined $first_results) {
    warn "Couldn't read any blast results from $file\n";
    unlink $tab_file, $crunch_file;
    # touch the files
    open F, ">$tab_file";
    close F;
    open F, ">$crunch_file";
    close F;
    return;
  }

  while (defined (my $hit = $first_results->next_hit)) {

    my $s_desc = $hit->s_desc;

    my $s_name = $hit->s_id;

    if (!defined $s_name) {
      $s_name = "unknown";
    }

    # don't put / in labels
    $s_name =~ s!/!-!g;

    my $hsp;
    while (defined ($hsp = $hit->next_hsp)) {
      my $q_start = $hsp->q_begin + $offset;
      my $q_end = $hsp->q_end + $offset;

      my $s_start = $hsp->s_begin;
      my $s_end = $hsp->s_end;

      my $percent_id = $hsp->percent;

      my $score = $hsp->score;

      if ($score < $min_score) {
        next;
      }

      if ($q_end < $q_start) {
        ($q_end, $q_start) = ($q_start, $q_end);
      }

      if ($s_end < $s_start) {
        ($s_end, $s_start) = ($s_start, $s_end);
      }

      printf CRUNCH_OUT "%d %3.2f %d %d %s %d %d %s %s\n",
        $hsp->score,
        $percent_id,
        (do {
          if ($hsp->q_strand == -1) {
            ($q_end, $q_start)
          } else {
            ($q_start, $q_end)
          }
        }),
        $file,
        (do {
          if (defined $hsp->s_strand and $hsp->s_strand == -1) {
            ($s_end, $s_start);
          } else {
            ($s_start, $s_end);
          }
        }),
        $s_name,
        $s_desc;

      if ($hsp->q_strand == -1) {
        print TAB_OUT "FT   \U$blast_program\E_HIT      complement($q_end..$q_start)\n";
      } else {
        print TAB_OUT "FT   \U$blast_program\E_HIT      $q_start..$q_end\n";
      }

      print TAB_OUT "FT                   /note=\"match to $s_name $s_start..$s_end  blast score $score  percent identity $percent_id  $s_desc\"\n";
      print TAB_OUT "FT                   /label=$s_name\n";
      print TAB_OUT "FT                   /score=$percent_id\n";

      my $colour = "255" . (" " . int((100 - $percent_id) * 2.55)) x 2;

      print TAB_OUT "FT                   /colour=$colour\n";
    }
  }

  close CRUNCH_OUT;
  close TAB_OUT;
}

sub plato_load
{
  if (open LSLOAD, "$lsload plato|") {
    my $first_line = <LSLOAD>;
    if (defined $first_line) {
      my $second_line = <LSLOAD>;

      if (defined $second_line && $second_line =~ /plato\s+\S+\s+([\d\.]+)/) {
        close LSLOAD;
        return $1;
      }
    }
  }
  warn "couldn't open pipe from lsload\n";
  close LSLOAD;
  return 9999;
}

sub start_child ($$)
{
  my ($file, $offset) = @_;

  my $output_file = "$file.out";

  my $child_pid = fork;

  if ($child_pid == 0) {
    # the child

    if (check_blast_output $output_file) {
      print STDERR "skipping $file - $output_file seems OK\n";
    } else {
      open(SAVEOUT, ">&STDOUT");

      # to avoid a warning about SAVEOUT being used only once
      print SAVEOUT "";

      close STDOUT;
      open STDOUT, ">$output_file" or die "can't open $output_file\n";

      my @exec_args = ();

      if ($use_blast2) {
        @exec_args = ($blast2_program,
                      "-d", $database,
                      "-i", $file, @blast_options);
      } else {
        @exec_args = ($blast_program, $database, $file, @blast_options);
      }

      if ($use_lsrun) {
        print STDERR "starting @exec_args\n";
        if ($use_plato && plato_load () < 13 && rand 1 < 0.5) {
          system "nice", "-5", $lsrun, "-m", "plato", "-v", @exec_args;
        } else {
          system "nice", "-5", $lsrun, "-v", @exec_args;
        }
      } else {
        print STDERR "starting @exec_args on localhost\n";
        system @exec_args;
      }

      close STDOUT;
      open(STDOUT, ">&SAVEOUT");
    }

    process_blast $output_file, $offset;
    exit (0);
  } else {
    if ($child_pid == -1) {
      die "$0: failed to start child process\n";
    } else {
      # the parent
      push @children, $child_pid;

      print STDERR "started process with id $child_pid\n";

      if ($use_lsrun) {
        # sleep so that the machines aren't swamped by jobs
        if (length $sequence > 100000 and length $sequence <= 1000000) {
          sleep 2;
        } else {
          if (length $sequence > 1000000) {
            sleep 3;
          } else {
            sleep 1;
          }
        }
      }
    }
  }
}

sub wait_for_child
{
  my $child_pid = wait;

  @children = grep { $_ != $child_pid } @children;

  print STDERR "$blast_program process with id $child_pid has finished\n";
}


for (my $i = 0 ; $i < @piece_files ; ++$i) {
  my $file = $piece_files[$i];

  if (@children < $child_count) {
    start_child $file, $piece_offsets[$i];
  } else {
    wait_for_child;
    start_child $file, $piece_offsets[$i];
  }
}

while (@children) {
  wait_for_child;
}


# wait for nfs
sleep 10 if $use_lsrun;

#  for (my $i = 0 ; $i < @piece_files ; ++$i) {
#    my $file = $piece_files[$i];

#    process_blast ($file . ".out", $piece_offsets[$i]);
#  }

open CRUNCH_OUT, ">$base_name.crunch";

for my $in_file (map {"$_.out.crunch"} @piece_files) {
  open IN_FILE, $in_file or die "can't open $in_file\n";
  print CRUNCH_OUT <IN_FILE>;
  close IN_FILE;
}

close CRUNCH_OUT;

print STDERR "wrote MSPcrunch output to $base_name.crunch\n";

open TAB_OUT, ">$base_name.tab";

for my $in_file (map {"$_.out.tab"} @piece_files) {
  open IN_FILE, $in_file or die "can't open $in_file\n";
  print TAB_OUT <IN_FILE>;
  close IN_FILE;
}

close TAB_OUT;

print STDERR "wrote tab file output to $base_name.tab\n";
