#!/usr/bin/perl -w

# cure_refseqs.pl
# A script to update refseqs (i.e. "cure" IHPs) using  consensus sequences obtained by
# alignments with sub-genome reads (e.g. 'A' genome reads vs. full 95k set derived from
# A, C and AACC Brassica genomes)
#
# Martin.Trick@bbsrc.ac.uk, 9/4/2010

use FileHandle;
STDOUT->autoflush(1);

# Hashes for original refseqs and consensus seqs (and quality strings - for future use)
my %ref = ();
my %cns = ();
my %qual = ();

# CNS Phred quality threshold required for change to refseq
my $threshold = 40;

# NB. maq cns2fq has default thresholds *for reads* of mapping quality >40 and depth >= 3
# Depth becomes encoded in the case of characters in sequence string

# Script requires the refseq fasta file, the fastq file of the consensus and writes an updated
# refseq fasta file. The accompanying wrapper script "cure_cycle.pl" controls these filenames
# Check arguments
unless (@ARGV == 3) {
  print "Usage: cure_refsqs.pl <input refseq fasta file> <input cns fastq file> <output fasta file>\n";
  exit;
}

# Open FHs
open (REF, "<$ARGV[0]") or die "Couldn't open input fasta file ($!)\n";
open (CNS, "<$ARGV[1]") or die "Couldn't open input fastq file ($!)\n";
open (OUT, ">$ARGV[2]") or die "Couldn't open output file ($!)\n";

# Read in refseqs
print "Reading refseqs... ";
while (<REF>) {
  chomp;
  if (/^>(\S+)(.*)/){
    ($id) = $1;
  }
  elsif (/^(\w+)$/){
    $ref{$id} .= $1;
  }
}
close REF;
print "done\n";

# Need to be a little cleverer to parse multiline fastq file
# Couldn't figure out how to do it more conventionally. so this ...
my $id;
my $quality;
print "Reading consensus sequences... ";
while (<CNS>) {
  chomp;
  if (/^\@(\w+)$/ && !defined $id && !$quality ){
    ($id) = $1;
    next;
  }
  if (/^(\w+)$/ && !$quality){ # \w will let apparent hemi-SNPs through
    $cns{$id} .= $1; # preserve case -  cns2fq's encoding of depth in sequence string
  }
  if (/^\+$/ && !$quality) {
    $quality = 1;
    next;
  }
  if ($quality) {
    $qual{$id} .= $_;
    if (length($qual{$id}) == length($cns{$id})) { # We've come to the end of this record
      undef $quality;
      undef $id;
    }
  }
}
close CNS;
print "done (" . scalar(keys(%cns)) . " records parsed)\n";

# Now cure refseqs
my $ambiguity = 0;
my $qual = 0;
my $cured = 0;
print "Curing refseqs... ";
foreach my $id (sort keys %cns) {
  # Skip this entire refseq if we have uniformly poor depth contributing to CNS
  next unless ($cns{$id} =~ /[^nacgt]/);
  # Else convert from string to array operations (expensive?)
  my @ref = split //, $ref{$id};
  my @cns = split //, $cns{$id};
  my @qual = split //, $qual{$id};
  foreach my $pos(0 .. $#ref) {
    next if (($cns[$pos] =~ /$ref[$pos]/i) || ($cns[$pos] =~ /[a-z]/)); # no difference or poor depth
    if (ord($qual[$pos]) - 33 >= $threshold) {
      $qual++;
      if ($cns[$pos] !~ /[acgt]/i) {
	$ambiguity++ ;
	#print "CNS for $id at pos " . ($pos+1) . " is ambiguous ($cns[$pos])\n";
      }
      else {
	$ref[$pos] = $cns[$pos];
	$cured++;
      }
    }
  }
  # Construct the cured refseq (NB, no EOLs in sequence string)
  $ref{$id} = join '', @ref;
}
print "done\n";
print "$ambiguity ambiguity codes (>= Phred $threshold) counted in consensus\n";
#print "" . ($ambiguity/$qual)*100 . "% of $qual refseq bases cured by consensus\n";
print "$cured bases cured by consensus\n";

foreach my $id (sort keys %ref) {
  print  OUT ">$id\n$ref{$id}\n";
}
close OUT;