#!/usr/bin/env perl

#
# Copyright 2011, Ben Langmead <langmea@cs.jhu.edu>
#
# This file is part of Bowtie 2.
#
# Bowtie 2 is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Bowtie 2 is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Bowtie 2.  If not, see <http://www.gnu.org/licenses/>.
#

# bowtie2:
#
# A wrapper script for bowtie2.  Provides various advantages over running
# bowtie2 directly, including:
#
# 1. Handling compressed inputs
# 2. Redirecting output to various files
# 3. Output directly to bam (not currently supported)

use strict;
use warnings;
use Getopt::Long;
use File::Spec;
use POSIX;

my ($vol,$script_path,$prog) 
                = File::Spec->splitpath(File::Spec->rel2abs( __FILE__ ));
my $os_is_nix   = ($^O eq "linux") || ($^O eq "darwin");
my $align_bin   = $os_is_nix ? 'bowtie2-align' : 'bowtie2-align.exe'; 
my $build_bin   = $os_is_nix ? 'bowtie2-build' : 'bowtie2-build.exe';               
my $align_prog  = File::Spec->catpath($vol,$script_path,$align_bin);
my $build_prog  = File::Spec->catpath($vol,$script_path,$build_bin);
my %signo       = ();
my @signame     = ();


{
	# Get signal info
	use Config;
	my $i = 0;
	for my $name (split(' ', $Config{sig_name})) {
		$signo{$name} = $i;
		$signame[$i] = $name;
		$i++;
	}
}

(-x "$align_prog") ||
	die "Error: Expected bowtie2 to be in same directory with bowtie2-align:\n$script_path";

# Get description of arguments from Bowtie 2 so that we can distinguish Bowtie
# 2 args from wrapper args
sub getBt2Desc($) {
	my $d = shift;
	my $cmd = "$align_prog --wrapper basic-0 --arg-desc";
	open(my $fh, "$cmd |") || die "Failed to run command '$cmd'";
	while(readline $fh) {
		chomp;
		next if /^\s*$/;
		my @ts = split(/\t/);
		$d->{$ts[0]} = $ts[1];
	}
	close($fh);
	$? == 0 || die;
}

my %desc = ();
my %wrapped = ("1" => 1, "2" => 1);
getBt2Desc(\%desc);

# Given an option like -1, determine whether it's wrapped (i.e. should be
# handled by this script rather than being passed along to Bowtie 2)
sub isWrapped($) { return defined($wrapped{$_[0]}); }

my @orig_argv = @ARGV;

my @bt2w_args = (); # options for wrapper
my @bt2_args  = (); # options for Bowtie 2
my $saw_dd = 0;
for(0..$#ARGV) {
	if($ARGV[$_] eq "--") {
		$saw_dd = 1;
		next;
	}
	push @bt2w_args, $ARGV[$_] if !$saw_dd;
	push @bt2_args,  $ARGV[$_] if  $saw_dd;
}
if(!$saw_dd) {
	@bt2_args = @bt2w_args;
	@bt2w_args= ();
}

my $debug = 0;
my %read_fns = ();
my %read_compress = ();
my $cap_out = undef;       # Filename for passthrough
my $no_unal = 0;

# Remove whitespace
for my $i (0..$#bt2_args) {
	$bt2_args[$i]=~ s/^\s+//; $bt2_args[$i] =~ s/\s+$//;
}

# We've handled arguments that the user has explicitly directed either to the
# wrapper or to bowtie2, now we capture some of the bowtie2 arguments that
# ought to be handled in the wrapper
for(my $i = 0; $i < scalar(@bt2_args); $i++) {
	next unless defined($bt2_args[$i]);
	my $arg = $bt2_args[$i];
	my @args = split(/=/, $arg);
	if(scalar(@args) > 2) {
		$args[1] = join("=", @args[1..$#args]);
	}
	$arg = $args[0];
	if($arg eq "-U" || $arg eq "--unpaired") {
		$bt2_args[$i] = undef;
		$arg =~ s/^-U//; $arg =~ s/^--unpaired//;
		if($arg ne "") {
			# Argument was part of this token
			my @args = split(/,/, $arg);
			for my $a (@args) { push @bt2w_args, ("-U", $a); }
		} else {
			# Argument is in the next token
			$i < scalar(@bt2_args)-1 || die;
			$i++;
			my @args = split(/,/, $bt2_args[$i]);
			for my $a (@args) { push @bt2w_args, ("-U", $a); }
			$bt2_args[$i] = undef;
		}
	}
	if($arg =~ /^--?([12])/ && $arg !~ /^--?12/) {
		my $mate = $1;
		$bt2_args[$i] = undef;
		$arg =~ s/^--?[12]//;
		if($arg ne "") {
			# Argument was part of this token
			my @args = split(/,/, $arg);
			for my $a (@args) { push @bt2w_args, ("-$mate", $a); }
		} else {
			# Argument is in the next token
			$i < scalar(@bt2_args)-1 || die;
			$i++;
			my @args = split(/,/, $bt2_args[$i]);
			for my $a (@args) { push @bt2w_args, ("-$mate", $a); }
			$bt2_args[$i] = undef;
		}
	}
	if($arg eq "--debug") {
		$debug = 1;
		$bt2_args[$i] = undef;
	}
	if($arg eq "--no-unal") {
		$no_unal = 1;
		$bt2_args[$i] = undef;
	}
	for my $rarg ("un-conc", "al-conc", "un", "al") {
		if($arg =~ /^--${rarg}$/ || $arg =~ /^--${rarg}-gz$/ || $arg =~ /^--${rarg}-bz2$/) {
			$bt2_args[$i] = undef;
			if(scalar(@args) > 1 && $args[1] ne "") {
				$read_fns{$rarg} = $args[1];
			} else {
				$i < scalar(@bt2_args)-1 || die "Error: --${rarg}* option takes an argument";
				$read_fns{$rarg} = $bt2_args[$i+1];
				$bt2_args[$i+1] = undef;
			}
			$read_compress{$rarg} = "";
			$read_compress{$rarg} = "gzip"  if $arg eq "--${rarg}-gz";
			$read_compress{$rarg} = "bzip2" if $arg eq "--${rarg}-bz2";
			last;
		}
	}
}
# If the user asked us to redirect some reads to files, or to suppress
# unaligned reads, then we need to capture the output from Bowtie 2 and pass it
# through this wrapper.
my $passthru = 0;
if(scalar(keys %read_fns) > 0 || $no_unal) {
	$passthru = 1;
	push @bt2_args, "--passthrough";
	$cap_out = "-";
	for(my $i = 0; $i < scalar(@bt2_args); $i++) {
		next unless defined($bt2_args[$i]);
		my $arg = $bt2_args[$i];
		if($arg eq "-S" || $arg eq "--output") {
			$i < scalar(@bt2_args)-1 || die "Error: -S/--output takes an argument";
			$cap_out = $bt2_args[$i+1];
			$bt2_args[$i] = undef;
			$bt2_args[$i+1] = undef;
		}
	}
}
my @tmp = ();
for (@bt2_args) { push(@tmp, $_) if defined($_); }
@bt2_args = @tmp;

my @unps = ();
my @mate1s = ();
my @mate2s = ();
my @to_delete = ();
my $temp_dir = "/tmp";
my $bam_out = 0;
my $ref_str = undef;
my $no_pipes = 0;
my $keep = 0;
my $verbose = 0;
my $readpipe = undef;

my @bt2w_args_cp = @bt2w_args;
@ARGV = @bt2w_args;
GetOptions(
	"1=s"                           => \@mate1s,
	"2=s"                           => \@mate2s,
	"reads=s"                       => \@unps,
	"U=s"                           => \@unps,
	"temp-directory=s"              => \$temp_dir,
	"bam"                           => \$bam_out,
	"no-named-pipes"                => \$no_pipes,
	"ref-string|reference-string=s" => \$ref_str,
	"keep"                          => \$keep,
	"verbose"                       => \$verbose
) || die "Bad option";

if($verbose) {
	print STDERR "Before arg handling:\n";
	print STDERR "  Wrapper args:\n[ @bt2w_args_cp ]\n";
	print STDERR "  Binary args:\n[ @bt2_args ]\n";
}

sub cat_file($$) {
	my ($ifn, $ofh) = @_;
	my $ifh = undef;
	if($ifn =~ /\.gz$/) {
		open($ifh, "gzip -dc $ifn |") ||
			die "Error: could not open gzipped read file: $ifn";
	} elsif($ifn =~ /\.bz2/) {
		open($ifh, "bzip2 -dc $ifn |") ||
			die "Error: could not open bzip2ed read file: $ifn";
	} else {
		open($ifh, $ifn) || die "Error: could not open read file: $ifn";
	}
	while(readline $ifh) { print {$ofh} $_; }
	close($ifh);
}

# Return non-zero if and only if the input should be wrapped (i.e. because
# it's compressed).
sub wrapInput($$$) {
	my ($unps, $mate1s, $mate2s) = @_;
	for my $fn (@$unps, @$mate1s, @$mate2s) {
		return 1 if $fn =~ /\.gz$/ || $fn =~ /\.bz2$/;
	}
	return 0;
}

if(wrapInput(\@unps, \@mate1s, \@mate2s)) {
	if(scalar(@mate2s) > 0) {
		#
		# Wrap paired-end inputs
		#
		# Put reads into temporary files or fork off processes to feed named pipes
		scalar(@mate2s) == scalar(@mate1s) ||
			die "Different number of files specified with --reads/-1 as with -2";
		# Make a named pipe for delivering mate #1s
		my $m1fn = "$temp_dir/$$.inpipe1";
		push @to_delete, $m1fn;
		push @bt2_args, "-1 $m1fn";
		# Create named pipe 1 for writing
		if(!$no_pipes) {
			mkfifo($m1fn, 0700) || die "Error: mkfifo($m1fn) failed.";
		}
		my $pid = 0;
		$pid = fork() unless $no_pipes;
		if($pid == 0) {
			# Open named pipe 1 for writing
			open(my $ofh, ">$m1fn") || die "Can't open '$m1fn' for writing";
			for my $ifn (@mate1s) { cat_file($ifn, $ofh); }
			close($ofh);
			exit 0 unless $no_pipes;
		}
		# Make a named pipe for delivering mate #2s
		my $m2fn = "$temp_dir/$$.inpipe2";
		push @to_delete, $m2fn;
		push @bt2_args, "-2 $m2fn";
		# Create named pipe 2 for writing
		if(!$no_pipes) {
			mkfifo($m2fn, 0700) || die "Error: mkfifo($m2fn) failed.";
		}
		$pid = 0;
		$pid = fork() unless $no_pipes;
		if($pid == 0) {
			# Open named pipe 2 for writing
			open(my $ofh, ">$m2fn") || die "Can't open '$m2fn' for writing";
			for my $ifn (@mate2s) { cat_file($ifn, $ofh); }
			close($ofh);
			exit 0 unless $no_pipes;
		}
	}
	if(scalar(@unps) > 0) {
		#
		# Wrap unpaired inputs.
		#
		# Make a named pipe for delivering unpaired reads
		my $ufn = "$temp_dir/$$.unp";
		push @to_delete, $ufn;
		push @bt2_args, "-U $ufn";
		# Create named pipe 2 for writing
		if(!$no_pipes) {
			mkfifo($ufn, 0700) || die "Error: mkfifo($ufn) failed.";
		}
		my $pid = 0;
		$pid = fork() unless $no_pipes;
		if($pid == 0) {
			# Open named pipe 2 for writing
			open(my $ofh, ">$ufn") || die "Can't open '$ufn' for writing";
			for my $ifn (@unps) { cat_file($ifn, $ofh); }
			close($ofh);
			exit 0 unless $no_pipes;
		}
	}
} else {
	if(scalar(@mate2s) > 0) {
		# Just pass all the mate arguments along to the binary
		push @bt2_args, ("-1", join(",", @mate1s));
		push @bt2_args, ("-2", join(",", @mate2s));
	}
	if(scalar(@unps) > 0) {
		push @bt2_args, ("-U", join(",", @unps));
	}
}

if(defined($ref_str)) {
	my $ofn = "$temp_dir/$$.ref_str.fa";
	open(my $ofh, ">$ofn") ||
		die "Error: could not open temporary fasta file '$ofn' for writing";
	print {$ofh} ">1\n$ref_str\n";
	close($ofh);
	push @to_delete, $ofn;
	system("$build_bin $ofn $ofn") == 0 ||
		die "Error: bowtie2-build returned non-0 exit level";
	push @bt2_args, ("--index", "$ofn");
	push @to_delete, ("$ofn.1.bt2", "$ofn.2.bt2", "$ofn.3.bt2", "$ofn.4.bt2",
	                  "$ofn.rev.1.bt2", "$ofn.rev.2.bt2");
}

if($verbose) {
	print STDERR "After arg handling:\n";
	print STDERR "  Binary args:\n[ @bt2_args ]\n";
}

my $debug_str = ($debug ? "-debug" : "");

# Construct command invoking bowtie2-align
my $cmd = "$align_prog$debug_str --wrapper basic-0 ".join(" ", @bt2_args);

# Possibly add read input on an anonymous pipe
$cmd = "$readpipe $cmd" if defined($readpipe);

print STDERR "$cmd\n" if $verbose;
my $ret;
if(defined($cap_out)) {
	# Open Bowtie 2 pipe
	open(BT, "$cmd |") || die "Error: Could not open Bowtie 2 pipe: '$cmd |'";
	# Open output pipe
	my $ofh = *STDOUT;
	my @fhs_to_close = ();
	if($cap_out ne "-") {
		open($ofh, ">$cap_out") ||
			die "Error: Could not open output file '$cap_out' for writing";
	}
	my %read_fhs = ();
	for my $i ("al", "un", "al-conc", "un-conc") {
		if(defined($read_fns{$i})) {
			if($i =~ /-conc$/) {
				# Open 2 output files, one for mate 1, one for mate 2
				my ($fn1, $fn2) = ($read_fns{$i}, $read_fns{$i});
				if($fn1 =~ /%/) {
					$fn1 =~ s/%/1/g; $fn2 =~ s/%/2/g;
				} elsif($fn1 =~ /\.[^.]*$/) {
					$fn1 =~ s/\.([^.]*)$/.1.$1/;
					$fn2 =~ s/\.([^.]*)$/.2.$1/;
				} else {
					$fn1 .= ".1";
					$fn2 .= ".2";
				}
				$fn1 ne $fn2 || die "$fn1\n$fn2\n";
				my ($redir1, $redir2) = (">$fn1", ">$fn2");
				$redir1 = "| gzip -c $redir1"  if $read_compress{$i} eq "gzip";
				$redir1 = "| bzip2 -c $redir1" if $read_compress{$i} eq "bzip2";
				$redir2 = "| gzip -c $redir2"  if $read_compress{$i} eq "gzip";
				$redir2 = "| bzip2 -c $redir2" if $read_compress{$i} eq "bzip2";
				open($read_fhs{$i}{1}, $redir1) || die "Error: Could not open --$i mate-1 output file '$fn1'";
				open($read_fhs{$i}{2}, $redir2) || die "Error: Could not open --$i mate-2 output file '$fn2'";
				push @fhs_to_close, $read_fhs{$i}{1};
				push @fhs_to_close, $read_fhs{$i}{2};
			} else {
				my $redir = ">$read_fns{$i}";
				$redir = "| gzip -c $redir"  if $read_compress{$i} eq "gzip";
				$redir = "| bzip2 -c $redir" if $read_compress{$i} eq "bzip2";
				open($read_fhs{$i}, $redir) || die "Error: Could not open --$i output file '$read_fns{$i}'";
				push @fhs_to_close, $read_fhs{$i};
			}
		}
	}
	while(<BT>) {
		chomp;
		my $filt = 0;
		unless(substr($_, 0, 1) eq "@") {
			# If we are supposed to output certain reads to files...
			my $tab1_i = index($_, "\t") + 1;
			my $tab2_i = index($_, "\t", $tab1_i);
			my $fl = substr($_, $tab1_i, $tab2_i - $tab1_i);
			my $unal = ($fl & 4) != 0;
			$filt = 1 if $no_unal && $unal;
			if($passthru) {
				if($filt) {
					# Next line is read with some whitespace escaped, which we
					# ignore b/c the record is filtered out by --no-unal
					my $l = <BT>;
				} else {
					my $mate1 = (($fl &  64) != 0);
					my $mate2 = (($fl & 128) != 0);
					my $unp = !$mate1 && !$mate2;
					my $pair = !$unp;
					# Next line is read with some whitespace escaped
					my $l = <BT>;
					chomp($l);
					$l =~ s/%(..)/chr(hex($1))/eg;
					if((defined($read_fhs{un}) || defined($read_fhs{al})) && $unp) {
						if($unal) {
							# Failed to align
							print {$read_fhs{un}} $l if defined($read_fhs{un});
						} else {
							# Aligned
							print {$read_fhs{al}} $l if defined($read_fhs{al});
						}
					}
					if((defined($read_fhs{"un-conc"}) || defined($read_fhs{"al-conc"})) && $pair) {
						my $conc  = (($fl &   2) != 0);
						if     ($conc && $mate1) {
							print {$read_fhs{"al-conc"}{1}} $l if defined($read_fhs{"al-conc"});
						} elsif($conc && $mate2) {
							print {$read_fhs{"al-conc"}{2}} $l if defined($read_fhs{"al-conc"});
						} elsif(!$conc && $mate1) {
							print {$read_fhs{"un-conc"}{1}} $l if defined($read_fhs{"un-conc"});
						} elsif(!$conc && $mate2) {
							print {$read_fhs{"un-conc"}{2}} $l if defined($read_fhs{"un-conc"});
						}
					}
				}
			}
		}
		print {$ofh} "$_\n" if !$filt;
	}
	for my $k (@fhs_to_close) { close($k); }
	close($ofh);
	close(BT);
	$ret = $?;
} else {
	$ret = system($cmd);
}
if(!$keep) { for(@to_delete) { unlink($_); } }

if ($ret == -1) {
    print STDERR "Failed to execute bowtie2-align: $!\n";
	exit 1;
} elsif ($ret & 127) {
	my $signm = "(unknown)";
	$signm = $signame[$ret & 127] if defined($signame[$ret & 127]);
	my $ad = "";
	$ad = "(core dumped)" if (($ret & 128) != 0);
    printf STDERR "bowtie2-align died with signal %d (%s) $ad\n", ($ret & 127), $signm;
	exit 1;
} elsif($ret != 0) {
	printf STDERR "bowtie2-align exited with value %d\n", ($ret >> 8);
}
exit ($ret >> 8);