#!/usr/bin/perl -w
#
# munge_barcoded_reads.pl
#
# Script to endoreduplicate the 6nt barcodes in tagged, multiplexed
# Illumina reads to allow Maq alignments to authentic pseudoref seqs (with
# their barcodes messed with in precisely the same way)
#
# Martin.Trick@bbsrc.ac.uk, 4/6/2009
#
# v2 - first look at real data revealed nibbling of barcodes, so let's
# work backwards from the M13 primer
#
# Martin.Trick@bbsrc.ac.uk, 28/8/2009

# v3 - Completely retooled for BAT TILLING 2 experiment with PE reads
# Martin.Trick@bbsrc.ac.uk, 23/3/2011
#
# Added code to discriminate between targets so that we can separately align
# and thus avoid homoeologue crosstalk
# Martin.Trick@bbsrc.ac.uk, 14/4/2011
#
# Think about disabling this option as we may we lose some signal doing
# it this way. Did it with a simple flag $sort. Set to zero, all munged reads
# should be sent to an arbitrary target file (first key in the hash).
#
#
# Found serious bug? Mates of M13-carrying reads were being munged x4 but
# they didn't have barcodes in the first place! Resulting in mixed read
# lengths - causing Maq misbehaviour? Adjust mate munging procedure by passing
# a $duplicate argument to the sub
# Martin.Trick@jic.ac.uk, 31/7/2012

use FileHandle;
STDOUT->autoflush(1);

unless (@ARGV >= 2) {
  print "Usage: mung_barcoded_reads_3.pl <fastq input file 1> <fastq input file 2> <barcode lookup file> [optional dump file]\n";
  exit;
}

# Messages
my $DEBUG = 0;
# Switch direct read-sorting on or off
my $sort = 0;

# Base positions discriminating target amplicons
my %targets = (
	       'FAD2_1_C5' => {'36' => 'C',
			       '59' => 'C',
			       '60' => 'C',
			      },
	       'FAD2_2_A5' => {'36' => 'T',
			       '59' => 'C',
			       '60' => 'C',
			    },
	       'FAD2_4_A1' => {'36' => 'C',
			       '59' => 'C',
			       '60' => 'T',
			      },
	       'GL2b_C6' => {'36' => 'T',
			     '59' => 'A',
			     '60' => 'C',
			    },
	      );

# String sort of hash keys to pick a filename for unsorted reads
my $default;
unless ($sort) {
  @output = sort keys %targets;
  $default = $output[0];
}

# Hash to hold FHs for output
my %filehandle = ();

# Open barcodes lookup
open (BARCODES, "<$ARGV[2]") or die "Couldn't open barcode lookup ($!)\n";
# Simple tabbed file - name of barcode and sequence (including M13 primer)

# Open output files, one for each target
foreach my $output(keys %targets) {
  open $filehandle{$output}, ">", "$output" . "_munged.fastq" or die "Couldn't open output file for $output ($!)";
}
# Optional file
open (DUMP, ">$ARGV[4]") if ($ARGV[4]); # Dump weird reads

# Parameters
my $duplicate = 4; # endoreduplication level for barcodes
my $barcode_length = 6;
my $nibbled = 4; # threshold for useful (reconstructable) remnant of barcode

# Various global counters, variables and flags
my $big_nibble = 0;
my $barcode;
my $full_code;
my $found_m13 = 0;
my $found_code = 0;
my $unique = 0;
my $stored = 0;

# Hashes for recognition of different amplicons
my %target_count = ();
my %spike_count = ();

# The anchor primer sequence
my $m13 = 'TGTAAAACGACGGCCAGT';

# The "spike" sequences
my %spikes = (
	      'JCVI_35759' => {'L' => 'GATCTCCGCTCCAAGATGG',
			       'R' => 'TGACGGTGAAGTACCTCCTC',
		    },
	      'JCVI_1358' => {'L' => 'GATGTCTTCAAGAGCCCAGC',
			      'R' => 'GATAGCAGTGACAATATCTCC',
		    },
	     );

# Hash for barcodes
my %barcodes = ();
# Hash to record surviving barcode lengths
my %nibbled = ();
# Hash to record found barcodes
my %found = ();
# Arrays of filehandles
my @fastq_fh = ();
my @index_fh = ();

print "Started at " . scalar localtime() . "\n";
# Capture barcodes - NB, unlike "mung_refs.pl", we are keying by sequence
while (<BARCODES>) {
  chomp;
  s/\r//g; # Windows FTP?
  my @cols = split;
  # Barcode isn't separated from the M13 primer in the list
  $cols[1] = substr($cols[1],0,$barcode_length);
  $barcodes{$cols[1]} = $cols[0];
}
close BARCODES;

# First thing to do is index the PE fastq files (if not already done)
print "Indexing fastq files... ";
foreach my $file($ARGV[0],$ARGV[1]) {
  local *FASTQ;
  local *INDEX;
  open (FASTQ, "<$file") or die "Couldn't open fastq file ($!)\n";
  # Don't clobber a pre-computed index file
  if (!-e "$file.idx" || -z "$file.idx") {
    open (INDEX, "+>$file.idx") or die "Couldn't open index file for read/write ($!)\n";
    build_index(*FASTQ, *INDEX);
  }
  else {
    open (INDEX, "<$file.idx");
  }
  # Store filehandles in array - these remain open
  push (@fastq_fh, *FASTQ);
  push (@index_fh, *INDEX);
}
print "done\n";

# Count lines in the fastq files
my $lines = `wc -l < $ARGV[0]`;
chomp ($lines);

# Go through the paired reads
print "Processing read-pairs... ";
READ_PAIR:for (my $line = 1; $line <= $lines; $line += 4) {
  @ids = ();
  @reads = ();
  @qualities = ();
  $munged = 0;
  $spike = 0;
  ++$read_pair;
  foreach $read(0 .. 1) {
    push @ids, get_line_with_index(\$fastq_fh[$read], \$index_fh[$read], $line);
    push @reads, get_line_with_index(\$fastq_fh[$read], \$index_fh[$read], ($line+1));
    push @qualities, get_line_with_index(\$fastq_fh[$read], \$index_fh[$read], ($line+3));
  }

  $stored++ if ($reads[0] ne '' && $reads[1] ne '');

 READ: foreach $read(0 .. 1) {
    next READ_PAIR if $munged;
    $mate = ($read == 0) ? 1 : 0;
    # Can we locate the M13 primer?
    if ($reads[$read] =~ m/$m13/g) {
      $found_m13++;
      $location = (pos $reads[$read]) - length($m13);
      $barcode = substr($reads[$read],0,$location);
      #next READ_PAIR unless (length($barcode) >= $nibbled);
      unless (length($barcode) >= $nibbled) {
	$big_nibble++;
	next READ_PAIR;
      }
      print "$ids[$read]\n$reads[$read]\n" if $DEBUG;
      print "Barcode $barcode\n" if $DEBUG;
      my $match = 0;
      my $full_code = '';
      # This code should select only reads with nibbled barcodes
      foreach $key(keys %barcodes) {
	if ($key =~ /$barcode$/) {
	  $match++;
	  $full_code = $key;
	  $found_code++ unless ($match > 1);
	  $nibbled{length($barcode)}++ unless ($match > 1);
	}
      }
      # Abandon this read-pair if we can't reconstruct the code
      next READ_PAIR unless ($match == 1);
      # This barcode remnant is unique, so we can reconstruct it
      $unique++;
      # Keep a record
      $found{$full_code}++;

      # Decide which target this read corresponds to (optional code?)
      my $cognate = ($sort) ? '' : $default; # Default target for unsorted reads
      if ($sort) {
      TARGET: foreach $target(keys %targets) {
	  foreach $pos(keys %{$targets{$target}}) {
	    next TARGET unless ($targets{$target}{$pos} eq substr($reads[$read],($location +$pos - 1),1));
	  }
	  $cognate = $target;
	  last TARGET;
	}
	next READ_PAIR unless (defined $cognate);
      }

      # Prepend trimmed read with munged full barcode
      &mung($read,$full_code,$location,$cognate);
      # We assume we can now mung the other read in this pair
      &mung($mate,$full_code,0,$cognate);
      $munged = 1;
      next READ_PAIR;
    } # M13
  } # READ
  unless ($munged) {
    # Else can we find a match to a spike?
  SPIKE: foreach my $spike(keys %spikes) {
      if (($reads[0] =~ m/$spikes{$spike}{L}/g) || ($reads[0] =~ m/$spikes{$spike}{R}/g)){
	$spike_count{$spike}++;
	$spike = 1;
	#last SPIKE;
	next READ_PAIR;
      }
    }
  }
  # Else we dump the unmunged read
  unless ($munged || $spike) {
    print DUMP "$ids[$read]\n$reads[$read]\n" if $ARGV[3];
    $dump++;
  }
} # READ_PAIR
print "done\n";


# Tidy up FHs
foreach $output(keys %targets) {
  close  $filehandle{$output};
}
close DUMP if $ARGV[3];


print "$read_pair read pairs processed\n\n";
print "$found_m13 primary reads recognised with perfect M13 match\n";
foreach $spike(sort keys %spike_count) {
  print "$spike_count{$spike} primary reads recognised with perfect match to spike $spike\n";
  $s+= $spike_count{$spike};
}
print "$big_nibble primary reads with wrecked barcodes < $nibbled nt\n";
print "$found_code primary reads with partial barcodes of >= $nibbled nt\n";
print "Barcode length distribution:\n";
foreach my $intact (sort {$a <=> $b} keys %nibbled){
  print "$intact- $nibbled{$intact}\n";
}
print "$unique reads with unique remnants of ". scalar(keys %found) . " barcodes\n";
print "\n";
print "". ($unique*2) . " reads munged OK\n";
print "" . ($found_m13 - $unique)*2 . " reads abandoned\n";
print "" . ($s*2) . " reads from spikes\n";
print "" . ($dump*2) . " unmatched reads\n";
print "" . ($read_pair*2) . " total reads\n";
print "\nFinished at " . scalar localtime() . "\n";


# Barcode and quality string munging, now as a subroutine
sub mung {
  my $read = shift;
  my $full_code = shift;
  my $location = shift;
  my $target = shift;
  print "Barcode remnant unique, reconstruct to $full_code\n" if $DEBUG;
  chop($ids[$read]);
  $ids[$read] .= "_$barcodes{$full_code}\n";
  $reads[$read] = substr($reads[$read],$location);
  substr($reads[$read],0,0) = ($full_code x= $duplicate);
  print "Printing munged read string $reads[$read]\n" if $DEBUG;
  print {$filehandle{$target}} "$ids[$read]$reads[$read]";
  # Mung the quality string, extrapolating missing barcode qualities from first
  $barcode_q = substr($qualities[$read],0,length($barcode));
  $first_q = substr($barcode_q,0,1);
  # Prepend with duplicate(s) of first q
  substr($barcode_q,0,0) = ($first_q x= ($barcode_length - length($barcode)));
  $barcode_q x= $duplicate;
  $qualities[$read] = substr($qualities[$read],$location);
  substr($qualities[$read],0,0) = $barcode_q;
  print "Munged quality string $qualities[$read]\n" if $DEBUG;
  print {$filehandle{$target}} "\+\n$qualities[$read]";
}

# Indexing and retrieval functions
sub build_index {
  my $fastq_file = shift;
  my $index_file = shift;
  my $offset = 0;

  while (<$fastq_file>) {
    print $index_file pack("Q", $offset);
    $offset = tell($fastq_file);
  }
}

sub get_line_with_index {
  my $fastq_file = shift;
  my $index_file = shift;
  my $line_number = shift;

  my $size;        # size of an index entry
  my $i_offset;    # offset into index of entry
  my $entry;       # index entry
  my $d_offset;    # offset into data file

  $size = length(pack("Q", 0));
  $i_offset = $size * ($line_number - 1);
  seek($index_file, $i_offset, 0) or return;
  read($index_file, $entry, $size);
  $d_offset = unpack("Q", $entry) or return;
  seek($fastq_file, $d_offset, 0) or return;
  return scalar(<$fastq_file>);
}


