
use strict;
use warnings;
use Getopt::Long;

use vars qw /$USAGE $IN $IN_F $OUT $OUT_F $BESTONLY/;

$USAGE = << "USE";
  usage: perl processMerged.pl -b [-i <file>] [-o <file>]
 
  options: -i      Path to SAM input file (default: stdin)
           -o      Path to SAM output file (default: stdout)
           -b      Enforce best-only.

  input:    Input can be provided by stdin or by a file in
            SAM format. It is assumed that reads have been
            mapped against different reference sequences using
            segemehl and the output files were merged using
            samtools or Picard tools with proper \@SQ entries.
            Furthermore, the SAM file MUST BE sorted by read
            names.

  output:   Output is in SAM format and is written to stdout or
            into a file. The perl skript updates the SAM tag 'NH'
            if necessary. If the best-only flag is set, only the
            best-scoring hits for each read (according to the SAM
            tag 'NM', lower values are better) are kept and the
            SAM tag 'NH' is updated to the number of best-scoring
            hits only.

 written by Christian Otto, Bioinformatics Leipzig, Mar 2012
            
USE

unless (GetOptions(
	    "i=s" => \$IN_F,
	    "o=s" => \$OUT_F,
	    "b" => \$BESTONLY)){
    printf STDERR $USAGE;
    exit -1;
}


if (!defined $IN_F && -t STDIN){
    printf STDERR $USAGE;
    exit -1;
}

sub main {
    my (@lines, $lastid, $lastedist, $lastscr, $line, $id, $scr);

    if (defined $IN_F){
	open($IN, "<".$IN_F) or die "Error in input\n";
    }
    else {
	$IN = \*STDIN;
    }
    
    if (defined $OUT_F){
	open($OUT, ">".$OUT_F) or die "Error in output\n";
    }
    else {
	$OUT = \*STDOUT;
    }

    $lastscr = -1;
    @lines = ();
    while(<$IN>){
	## ignore header lines
	do{print $OUT $_; next;} if /^@/;

	chomp;	
	$line = $_;
	$id = (split("\t"))[0];
	$scr = $1 if /NM:i:(\d+)/ or die;

	if (defined $lastid && $lastid eq $id){
	    if (defined $BESTONLY){
		## better than previous best
		if ($lastscr > $scr){
		    @lines = ();
		    $lastscr = $scr;
		}
		## worse than previous best
		elsif ($lastscr < $scr){
		    next;
		}
	    }
	    push(@lines, $line);
	}
	else {	
	    ## output    
	    print $OUT process(\@lines) if defined $lastid;

	    @lines = ($line);
	    $lastscr = $scr;
	    $lastid = $id;
	}
    }
    ## output
    print $OUT process(\@lines) if defined $lastid;

    ## close streams
    if (defined $IN_F){
	close($IN);
    }

    if (defined $OUT_F){
	close($OUT);
    }
}

sub process {
    my ($lines_ref) = @_;
    my ($i, $j, @F, $ret, $length);

    $ret = "";
    $length = scalar @{$lines_ref};

    for ($i = 0; $i < $length; $i++){
	@F = split("\t", $lines_ref->[$i]);
	for ($j = 11; $j < scalar(@F); $j++){
	    if ($F[$j] =~ m/NH:i:\d+/){
		$F[$j] = "NH:i:".$length;
		last;
	    }
	}
	die if $j == scalar(@F);
	$ret .= join("\t", @F)."\n";
    }    
    return ($ret);
}

main();
