#!/usr/bin/env perl

use strict;
use warnings;
use Getopt::Std;

my %opts = (t=>1);
getopts("SadskHo:R:x:t:", \%opts);

die('
Usage:   run-bwamem [options] <idxbase> <file1> [file2]

Options: -o STR    prefix for output files                       [inferred from input]
         -R STR    read group header line such as \'@RG\tID:foo\tSM:bar\'         [null]
         -x STR    read type: pacbio, ont2d or intractg                      [default]
                   intractg: intra-species contig (kb query, highly similar)
                   pacbio:   pacbio subreads (~10kb query, high error rate)
                   ont2d:    Oxford Nanopore reads (~10kb query, higher error rate)
         -t INT    number of threads                                               [1]

         -H        apply HLA typing
         -a        trim HiSeq2000/2500 PE resequencing adapters (via trimadap)
         -d        mark duplicate (via samblaster)
         -S        for BAM input, don\'t shuffle
         -s        sort the output alignment (via samtools; requring more RAM)
         -k        keep temporary files generated by typeHLA

Examples:

 * Map paired-end reads to GRCh38+ALT+decoy+HLA and perform HLA typing:

     run-bwamem -o prefix -t8 -HR"@RG\tID:foo\tSM:bar" hs38DH.fa read1.fq.gz read2.fq.gz

   Note: HLA typing is only effective for high-coverage data. The typing accuracy varies
   with the quality of input. It is only intended for research purpose, not for diagnostic.

 * Remap coordinate-sorted BAM, transfer read groups tags, trim Illumina PE adapters and
   sort the output. The BAM may contain single-end or paired-end reads, or a mixture of
   the two types. Specifying -R stops read group transfer.

     run-bwamem -sao prefix hs38DH.fa old-srt.bam

   Note: the adaptor trimmer included in bwa.kit is chosen because it fits the current
   mapping pipeline better. It is conservative and suboptimal. A more sophisticated
   trimmer is recommended if this becomes a concern.

 * Remap name-grouped BAM and mark duplicates:
 
     run-bwamem -Sdo prefix hs38DH.fa old-unsrt.bam

   Note: streamed duplicate marking requires all reads from a single paired-end library
   to be aligned at the same time.

Output files:

  {-o}.aln.bam - final alignment
  {-o}.hla.top - best genotypes for the 6 classical HLA genes (if there are HLA-* contigs)
  {-o}.hla.all - additional HLA genotypes consistent with data
  {-o}.log.*   - log files

') if @ARGV < 2;

my $idx = $ARGV[0];

my $exepath = $0 =~/^\S+\/[^\/\s]+/? $0 : &which($0);
my $root = $0 =~/^(\S+)\/[^\/\s]+/? $1 : undef;
$root = $exepath =~/^(\S+)\/[^\/\s]+/? $1 : undef if !defined($root);
die "ERROR: failed to locate the 'bwa.kit' directory\n" if !defined($root);

die("ERROR: failed to locate the BWA index. Please run '$root/bwa index -p $idx ref.fa'.\n")
  unless (-f "$idx.bwt" && -f "$idx.pac" && -f "$idx.sa" && -f "$idx.ann" && -f "$idx.amb");

if (@ARGV >= 3 && $ARGV[1] =~ /\.(bam|sam|sam\.gz)$/) {
	warn("WARNING: for SAM/BAM input, only the first sequence file is used.\n");
	@ARGV = 2;
}

if (defined($opts{p}) && @ARGV >= 3) {
	warn("WARNING: option -P is ignored as there are two input sequence files.\n");
	delete $opts{p};
}

my $prefix;
if (defined $opts{o}) {
	$prefix = $opts{o};
} elsif (@ARGV >= 3) {
	my $len = length($ARGV[1]) < length($ARGV[2])? length($ARGV[1]) : length($ARGV[2]);
	my $i;
	for ($i = 0; $i < $len; ++$i) {
		last if substr($ARGV[1], $i, 1) ne substr($ARGV[2], $i, 1)
	}
	$prefix = substr($ARGV[1], 0, $i) if $i > 0;
} elsif ($ARGV[1] =~ /^(\S+)\.(fastq|fq|fasta|fa|mag|mag\.gz|fasta\.gz|fa\.gz|fastq\.gz|fq\.gz|bam)$/) {
	$prefix = $1;
}
die("ERROR: failed to identify the prefix for output. Please specify -o.\n") unless defined($prefix);

my $size = 0;
my $comp_ratio = 3.;
for my $f (@ARGV[1..$#ARGV]) {
	my @a = stat($f);
	my $s = $a[7];
	die("ERROR: failed to read file $f\n") if !defined($s);
	$s *= $comp_ratio if $f =~ /\.(gz|bam)$/;
	$size += int($s) + 1;
}

my $is_pe = (defined($opts{p}) || @ARGV >= 3)? 1 : 0;
my $is_bam = $ARGV[1] =~ /\.bam$/? 1 : 0;

if (defined($opts{x})) {
	delete($opts{d}); delete($opts{a}); delete $opts{p};
}

# for BAM input, find @RG header lines
my @RG_lines = ();
if ($is_bam && !defined($opts{R})) {
	my $fh;
	open($fh, "$root/samtools view -H $ARGV[1] |") || die;
	while (<$fh>) {
		chomp;
		if (/^\@RG\t/) {
			s/\t/\\t/g;
			push(@RG_lines, "-H'$_'");
		}
	}
	close($fh);
}

warn("WARNING: many programs require read groups. Please specify with -R if you can.\n") if !defined($opts{R}) && @RG_lines == 0;

my $cmd = '';
if ($is_bam) {
	my $cmd_sam2bam = "cat $ARGV[1] \\\n";
	my $ntmps = int($size / 4e9) + 1;
	my $cmd_shuf = !defined($opts{S})? "  | $root/htsbox bamshuf -uOn$ntmps - $prefix.shuf \\\n" : "";
	my $bam2fq_opt = @RG_lines > 0? " -t" : "";
	my $cmd_bam2fq = "  | $root/htsbox bam2fq -O$bam2fq_opt - \\\n";
	$cmd = $cmd_sam2bam . $cmd_shuf . $cmd_bam2fq;
} elsif (@ARGV >= 3) {
	$cmd = "$root/seqtk mergepe $ARGV[1] $ARGV[2] \\\n";
} else {
	$cmd = "cat $ARGV[1] \\\n";
}

my $bwa_opts = "-p " . ($opts{t} > 1? "-t$opts{t} " : "") . (defined($opts{x})? "-x $opts{x} " : "") . (defined($opts{R})? "-R'$opts{R}' " : "");
$bwa_opts .= join(" ", @RG_lines) . " -C " if @RG_lines > 0;

$cmd .= "  | $root/trimadap 2> $prefix.log.trim \\\n" if defined($opts{a});
$cmd .= "  | $root/bwa mem $bwa_opts$ARGV[0] - 2> $prefix.log.bwamem \\\n";
$cmd .= "  | $root/samblaster 2> $prefix.log.dedup \\\n" if defined($opts{d});

my $has_hla = 0;
if (-f "$ARGV[0].alt") {
	my $fh;
	open($fh, "$ARGV[0].alt") || die;
	while (<$fh>) {
		$has_hla = 1 if /^HLA-[^\s\*]+\*\d+/;
	}
	close($fh);
	my $hla_pre = $has_hla? "-p $prefix.hla " : "";
	$cmd .= "  | $root/k8 $root/bwa-postalt.js $hla_pre$ARGV[0].alt \\\n";
}

my $t_sort = $opts{t} < 4? $opts{t} : 4;
$cmd .= defined($opts{s})? "  | $root/samtools sort -@ $t_sort -m1G - $prefix.aln;\n" : "  | $root/samtools view -1 - > $prefix.aln.bam;\n";

if ($has_hla && defined($opts{H}) && (!defined($opts{x}) || $opts{x} eq 'intractg')) {
	$cmd .= "$root/run-HLA ". (defined($opts{x}) && $opts{x} eq 'intractg'? "-A " : "") . "$prefix.hla > $prefix.hla.top 2> $prefix.log.hla;\n";
	$cmd .= "touch $prefix.hla.HLA-dummy.gt; cat $prefix.hla.HLA*.gt | grep ^GT | cut -f2- > $prefix.hla.all;\n";
	$cmd .= "rm -f $prefix.hla.HLA*;\n" unless defined($opts{k});
}

print $cmd;

sub which
{
	my $file = shift;
	my $path = (@_)? shift : $ENV{PATH};
	return if (!defined($path));
	foreach my $x (split(":", $path)) {
		$x =~ s/\/$//;
		return "$x/$file" if (-x "$x/$file");
	}
	return;
}