#!/usr/bin/env perl

# Copyright 2013-2015, Derrick Wood, Jennifer Lu <jlu26@jhmi.edu>
#
# This file is part of the Kraken taxonomic sequence classification system.
#
# Kraken 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.
#
# Kraken 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 Kraken.  If not, see <http://www.gnu.org/licenses/>.

# Merges two files specified on the command line, print to stdout
# Designed to be called by kraken

use strict;
use warnings;
use File::Basename;
use Getopt::Long;

my $PROG = basename $0;

my $fasta_input = 0;
my $fastq_input = 0;
my $gunzip = 0;
my $bunzip2 = 0;
my $check_names = 0;
my $output_format;

GetOptions(
  "fa" => \$fasta_input,
  "fq" => \$fastq_input,
  "gz" => \$gunzip,
  "bz2" => \$bunzip2,
  "check-names" => \$check_names,
  "output-format=s" => \$output_format
);

if (@ARGV != 2) {
  die "$PROG: must have exactly two filename arguments\n";
}
for my $file (@ARGV) {
  if (! -e $file) {
    die "$PROG: $file does not exist\n";
  } 
  if (! -f $file) {
    die "$PROG: $file is not a regular file\n";
  }
}

my $compressed = $gunzip || $bunzip2;
if ($gunzip && $bunzip2) {
  die "$PROG: can't use both gzip and bzip2 compression flags\n";
}
if (! ($fasta_input xor $fastq_input)) {
  die "$PROG: must specify exactly one of FASTA and FASTQ input flags\n";
}

# Open files (or pipes from decompression progs)
my ($fh1, $fh2);
if ($gunzip) {
  open $fh1, "-|", "gunzip", "-c", $ARGV[0]
    or die "$PROG: can't open gunzip pipe with $ARGV[0]: $!\n";
  open $fh2, "-|", "gunzip", "-c", $ARGV[1]
    or die "$PROG: can't open gunzip pipe with $ARGV[1]: $!\n";
}
elsif ($bunzip2) {
  open $fh1, "-|", "bunzip2", "-c", $ARGV[0]
    or die "$PROG: can't open bunzip2 pipe with $ARGV[0]: $!\n";
  open $fh2, "-|", "bunzip2", "-c", $ARGV[1]
    or die "$PROG: can't open bunzip2 pipe with $ARGV[1]: $!\n";
}
else {
  open $fh1, "<", $ARGV[0]
    or die "$PROG: can't open $ARGV[0]: $!\n";
  open $fh2, "<", $ARGV[1]
    or die "$PROG: can't open $ARGV[1]: $!\n";
}

# read/merge/print loop
# make sure names match before merging
my ($seq1, $seq2, $delimiter);
if ($output_format eq "legacy") {
  $delimiter = 'N';
}
else {
  $delimiter = '|';
}
while (defined($seq1 = read_sequence($fh1))) {
  $seq2 = read_sequence($fh2);
  if (! defined $seq2) {
    die "$PROG: mismatched sequence counts\n";
  }
  if ($check_names) {
    my $comparison_id1 = $seq1->{id} =~ /(\S+)/; # Only check up until first whitespace character
    my $comparison_id2 = $seq2->{id} =~ /(\S+)/;
    if ($comparison_id1 ne $comparison_id2) {
        die "$PROG: mismatched mate pair names ('$seq1->{id}' & '$seq2->{id}')\n";
    }
  }
  if ($fastq_input) {
    print_merged_sequence_fastq($seq1, $seq2, $delimiter);
  }
  else {
    print_merged_sequence($seq1, $seq2, $delimiter);
  }
}
if (defined($seq2 = read_sequence($fh2))) {
  die "$PROG: mismatched sequence counts\n";
}
close $fh1;
close $fh2;

{
  my %buffers;  # Needed due to fasta formatting
  sub read_sequence {
    my $fh = shift;
    my $id;
    my $seq = "";
    my $qual = "";
    if (! exists $buffers{$fh}) {
      $buffers{$fh} = <$fh>;
    }
    if (! defined $buffers{$fh}) {  # No more file
      return undef;
    }
    if ($fasta_input) {
      if ($buffers{$fh} =~ /^>(\S+)/) {
        $id = $1;
      }
      else {
        die "$PROG: malformed fasta file (line = '$buffers{$fh}')\n";
      }
      delete $buffers{$fh};
      while (<$fh>) {
        if (/^>/) {
          $buffers{$fh} = $_;
          last;
        }
        else {
          chomp;
          $seq .= $_;
        }
      }
    }
    elsif ($fastq_input) {
      my $fastq_header_regex;
      if ($output_format eq "legacy") {
	  $fastq_header_regex = qr/^@(\S+)/;
      }
      else {
	  $fastq_header_regex = qr/^@(\S+\s\S+)/; # Allow one whitespace character in fastq header
      }
      if ($buffers{$fh} =~ $fastq_header_regex) {  
        $id = $1;
      }
      else {
        if ($buffers{$fh} =~ /^$/) {
          return undef;  # Some fastq files end with blank line, allow it
        }
        die "$PROG: malformed fastq file (line = '$buffers{$fh}')\n";
      }
      delete $buffers{$fh};
      chomp($seq = <$fh>);
      scalar <$fh>;  # quality header
      chomp($qual = <$fh>); # quality values
    }
    else {
      # should never get here
      die "$PROG: I have no idea what kind of input I'm reading!!!\n";
    }

    return { id => $id, seq => $seq, qual => $qual };
  }
}

sub print_merged_sequence {
  my ($seq1, $seq2, $delimiter) = @_;
  print ">" . $seq1->{id} . "\n";
  print $seq1->{seq} . $delimiter . $seq2->{seq} . "\n";
}

sub print_merged_sequence_fastq {
  my ($seq1, $seq2, $delimiter) = @_;
  if ($output_format eq "legacy") {
    print "@" . $seq1->{id} . "\n";
  }
  else {
    print "@" . $seq1->{id} . $delimiter . $seq2->{id} . "\n";
  }
  print $seq1->{seq} . $delimiter . $seq2->{seq} . "\n";
  print "+\n";
  print $seq1->{qual} . $delimiter . $seq2->{qual} . "\n";
}
