#!/usr/common/usg/languages/perl/5.16.0/bin/perl -w
#
# scaffReportToFasta.pl by Jarrod Chapman <jchapman@lbl.gov> Thu Jun  2 15:08:06 PDT 2011
# Copyright 2011 Jarrod Chapman. All rights reserved.
#

use Getopt::Std;
my %opts = ();
my $validLine = getopts('c:s:g:G:', \%opts);
my @required = ("c","s",);
my $nRequired = 0;
map {$nRequired += exists($opts{$_})} @required;
$validLine &= ($nRequired == @required);
if ($validLine != 1) {
    print "Usage: ./scaffReportToFasta.pl <-c contigFastaFile> <-s scaffoldReportFile> <<-g gapClosureFileGlob>> <<-G minGapNs>>\n";
    exit;
}

my $date = `date`;
chomp $date;
print STDERR "$date\nscaffReportToFasta";
while (my ($opt,$val) = each(%opts)) {
    print STDERR " -$opt=$val";
}
print STDERR "\n";

my $contigFastaFile = $opts{"c"};
my $gapClosureFileGlob = undef;
if (exists($opts{"g"})) {
   $gapClosureFileGlob = $opts{"g"};
}
my $scaffReportFile = $opts{"s"};

my $minGapSize = 10;
if (exists($opts{"G"})) {
    $minGapSize = $opts{"G"};
}

print STDERR "Reading $contigFastaFile...\n";

my %contigSequence = ();
my %contigLengths = ();
my $currentEntry = undef;
my $seq = "";
my $seqLength = length($seq);
open (F,$contigFastaFile) || die "Couldn't open $contigFastaFile\n";
while (my $line = <F>) {
    chomp $line;
    if ($line =~ /^>/) {
	$seqLength = length($seq);
	if (defined($currentEntry)) {
	    $contigSequence{$currentEntry} = $seq;
	    $contigLengths{$currentEntry} = $seqLength;
	}
	$seq = "";
	($currentEntry) = $line =~ />(\S+)/;
    } else {
	$seq .= $line;
    }
}
close F;

$seqLength = length($seq);
if (defined($currentEntry)) {
    $contigSequence{$currentEntry} = $seq;
    $contigLengths{$currentEntry} = $seqLength;
}

print STDERR "Done.\n";

print STDERR "Reading $scaffReportFile...\n";

my @scaffoldIDs = ();
my %scaffoldInfo = ();

open (S,$scaffReportFile) || die "Couldn't open $scaffReportFile\n";
while (my $line = <S>) {
    chomp $line;
    my @cols = split(/\t/,$line);
    my $scaffID = $cols[0];
    unless (exists ($scaffoldInfo{$scaffID})) {
	push(@scaffoldIDs,$scaffID);
	$scaffoldInfo{$scaffID} = [];
    }
    if ($cols[1] =~ /^CONTIG/) {
	push(@{$scaffoldInfo{$scaffID}},"$cols[2],$cols[3],$cols[4]");
    }
}
close S;

print STDERR "Done.\n";

my %gapClosures = ();

if (defined($gapClosureFileGlob)) {
    
    my @gapClosureFiles = glob($gapClosureFileGlob);
    my $nGapClosureFiles = scalar(@gapClosureFiles);
    my $counter = 0;

    foreach my $gapClosureFile (@gapClosureFiles) {
	$counter++;

	print STDERR "Reading $gapClosureFile ($counter/$nGapClosureFiles)...\n";
	open (G,$gapClosureFile) || die "Couldn't open $gapClosureFile\n";
	while (my $line = <G>) {
	    chomp $line;
	    my ($scaffold,$contig1,$primer1,$contig2,$primer2,$closure) = split(/\t/,$line);

	    my ($c1) = $contig1 =~ /(.+)\.[35]$/; 
	    my ($c2) = $contig2 =~ /(.+)\.[35]$/; 
	    $gapClosures{"$c1:$c2"} = "$primer1,$primer2,$closure";

	}
	close G;
    }
}

my %printingScaffold = ();
my $lineWidth = 50;
my $printBuffer = 100*$lineWidth;
#my $scaffTotal = 0;

foreach my $scaffID (@scaffoldIDs) {
    my $scaffSeq = "";
    my $scaffCoord = 0;
    my $scaffLen = 0;
    my @contigList = @{$scaffoldInfo{$scaffID}};
    my $lastContigName = "NULL";
    foreach my $contigInfo (@contigList) {
	my ($cID,$sStart,$sEnd) = split(/\,/,$contigInfo);
	my ($cStrand,$cName) = $cID =~ /^([\+\-])(.+)$/;

	my $cSeq = $contigSequence{$cName};
	if ($cStrand eq "-") {
	    my $rc = reverse($cSeq);
	    $rc =~ tr/acgtACGT/tgcaTGCA/;
	    $cSeq = $rc;
	}

	my $gapName = "$lastContigName:$cName";
	my $reverseGapName = "$cName:$lastContigName";
	$lastContigName = $cName;

	if (exists($gapClosures{$gapName})) {
	    my $closureInfo = $gapClosures{$gapName};
	    my ($primer1,$primer2,$closure) = split(/\,/,$closureInfo);

	    my $p1Len = length($primer1);
	    my $scaffTail = uc(substr($scaffSeq,-$p1Len));
	    my $p2Len = length($primer2);
	    my $contigHead = uc(substr($cSeq,0,$p2Len));

	    unless (($scaffTail eq $primer1) && ($contigHead eq $primer2)) {
		die "Invalid gap closure for $gapName ($closureInfo)\n$scaffTail\n$primer1\n$primer2\n$contigHead\n";
	    }

	    my $prefix = substr($scaffSeq,0,$scaffLen-$p1Len);
	    my ($suffix) = $cSeq =~ /^.{$p2Len}(.*)$/;

	    $scaffSeq = $prefix.$closure.$suffix;

	} elsif (exists($gapClosures{$reverseGapName})) {

	    my $closureInfo = $gapClosures{$reverseGapName};
	    my ($primer2,$primer1,$closure) = split(/\,/,$closureInfo);
	    $primer2 = reverse($primer2);
	    $primer2 =~ tr/acgtACGT/tgcaTGCA/;
	    $primer1 = reverse($primer1);
	    $primer1 =~ tr/acgtACGT/tgcaTGCA/;
	    $closure = reverse($closure);
	    $closure =~ tr/acgtACGT/tgcaTGCA/;

	    my $p1Len = length($primer1);
	    my $scaffTail = uc(substr($scaffSeq,-$p1Len));
	    my $p2Len = length($primer2);
	    my $contigHead = uc(substr($cSeq,0,$p2Len));

	    unless (($scaffTail eq $primer1) && ($contigHead eq $primer2)) {
		die "Invalid gap closure for $reverseGapName ($closureInfo)\n$scaffTail\n$primer1\n$primer2\n$contigHead\n";
	    }

	    my $prefix = substr($scaffSeq,0,$scaffLen-$p1Len);
	    my ($suffix) = $cSeq =~ /^.{$p2Len}(.*)$/;

	    $scaffSeq = $prefix.$closure.$suffix;

	} elsif ($scaffSeq) {
	    my $gapSize = ($sStart - $scaffCoord - 1);
	    
	    my $gapSeq = "";
	    if ($gapSize < $minGapSize) {
		$gapSeq = "N"x$minGapSize;
	    } else {
		$gapSeq = "N"x$gapSize;
	    }
	    $scaffSeq .= $gapSeq.$cSeq;

	} else {
	    $scaffSeq = $cSeq;
	}
	$scaffCoord = $sEnd;
	$scaffLen = length($scaffSeq);
	if ($scaffLen > 2*$printBuffer) {
	    my $nBuffers = sprintf("%d",$scaffLen/$printBuffer);
	    my $printBlock = ($nBuffers-1)*$printBuffer;
	    my $printThis = substr($scaffSeq,0,$printBlock);
	    my $saveThis = substr($scaffSeq,$printBlock);
	    printFasta($scaffID,$printThis);
	    $scaffSeq = $saveThis;
	    $scaffLen = length($scaffSeq);
	}
    }
    printFasta($scaffID,$scaffSeq);

#    if ($scaffTotal > 1e6) {
#	last;
#    }
}


sub printFasta {
    my ($name,$seq) = @_;
    my $seqLen = length($seq);
#    $scaffTotal += $seqLen;

    my $bpl = $lineWidth;
    my $nLines = sprintf("%d",$seqLen/$bpl);
    if ($seqLen%$bpl != 0) {
        $nLines++;
    }

    unless (exists($printingScaffold{$name})) {
	print ">$name\n";
	$printingScaffold{$name} = 1;
    }
    for (my $j = 0;$j<$nLines;$j++) {
        my $seqLine = substr($seq,$j*$bpl,$bpl);
        my $text = $seqLine;
        $text.= "\n";
        print $text;
    }
}
