#!/usr/common/usg/languages/perl/5.16.0/bin/perl -w
#
# gapDivider.pl by Jarrod Chapman <jchapman@lbl.gov> Mon Jun 11 09:59:22 PDT 2012
# Copyright 2012 Jarrod Chapman. All rights reserved.
#

use Getopt::Std;
my %opts = ();
my $validLine = getopts('g:n:o:s:Q:M:', \%opts);
my @required = ("g","n","o","s");
my $nRequired = 0;
map {$nRequired += exists($opts{$_})} @required;
$validLine &= ($nRequired == @required);
if ($validLine != 1) {
    print "Usage: ./gapDivider.pl <-o outputPrefix> <-n gapsPerFile> <-g gapInfoFile> <-s srfFile> <<-Q outputQualityOffset (default is 33)>> <<-M modulusInfo (e.g. 3:10)>>\n";
    exit;
}

my $modulus = undef;
my $keepMod = undef;
if (exists($opts{"M"})) {
    my $modInfo = $opts{"M"};
    ($keepMod,$modulus) = $modInfo =~ /^(\d+):(\d+)$/;
}

my $defaultQualOffset = 33;
if (exists($opts{"Q"})) {
    $defaultQualOffset = $opts{"Q"};
}
if ($defaultQualOffset =~ /\D/) {
    die "Quality offset ($defaultQualOffset) should be integral (usually 33 or 64).\n";
}

my $gapsPerFile = $opts{"n"};
my $gapInfoFile = $opts{"g"};
my $srfFile = $opts{"s"};
my $outputPrefix = $opts{"o"};

my @gapDataFiles = ();
my @gapDataQuals = ();
my $nGapDataFiles  = 0;
open (G,$gapInfoFile) || die "Couldn't open $gapInfoFile\n";
while (my $line = <G>) {
    chomp $line;
    my ($file,$qualOffset) = split(/\t/,$line);
    push(@gapDataFiles,$file);
    push(@gapDataQuals,$qualOffset);
    $nGapDataFiles++;

}
close G;

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

my %gapIndex = ();
my $gapID = 0;
open (S,$srfFile) || die "Couldn't open $srfFile\n";
my $currentScaffold = "NULL";
my $currentGap = "";
while (my $line = <S>) {
    chomp $line;
    my (@cols) = split(/\t/,$line);
    my ($scaff,$label,$info) = @cols[0,1,2];
    unless ($label =~ /^CONTIG/) {
	next;
    }
    my ($contig) = $info =~ /^[\+\-](.+)$/;
    if ($scaff eq $currentScaffold) {
	$currentGap .= "$contig";
	$gapIndex{$currentGap} = $gapID;
	$gapID++;
    } else {
	$currentScaffold = $scaff;
    }
    $currentGap = "$scaff.$contig.";
}
close S;

print STDERR "Done.  Indexed $gapID gaps.\n";

my $gapBlockIndex = 0;
my $nGapBlocks = sprintf("%d",1+$gapID/$gapsPerFile);
my %fileHandles = ();
my @overflow = ();
while ($gapBlockIndex < $nGapBlocks) {

    my $minGapIndex = $gapBlockIndex*$gapsPerFile;
    my $maxGapIndex = (1+$gapBlockIndex)*$gapsPerFile - 1;

    my %gapData = ();

    my $goodIndex = 1;
    if (defined($modulus) && (($gapBlockIndex%$modulus) != $keepMod)) {
	$goodIndex = 0;
    }

    my $outFile = $outputPrefix.$gapBlockIndex;
    if ($goodIndex) {
	print STDERR "Writing $outFile...\n";
	open (OUT,">$outFile") || die "Couldn't open $outFile\n";
    }

    my $nReadsDumped = 0;

    my @spillOver = ();
    foreach my $overflew (@overflow) {
	my ($scaff,$c1,$p1,$c2,$p2,$gapSize,$gapUnc,@reads) = split(/\t/,$overflew);
	my ($contig1) = $c1 =~ /^(.+)\.[35]$/;
	my ($contig2) = $c2 =~ /^(.+)\.[35]$/;
	my $gapName = "$scaff.$contig1.$contig2";
	my $gapNumber = $gapIndex{$gapName};
	if ($gapNumber > $maxGapIndex) {
	    push(@spillOver,$overflew);
	} else {
	    if (exists($gapData{$gapNumber})) {
		push(@{$gapData{$gapNumber}},@reads);
	    } else {
		$gapData{$gapNumber} = [$scaff,$c1,$p1,$c2,$p2,$gapSize,$gapUnc,@reads];
	    }
	}
    }
    @overflow = @spillOver;

    for (my $fileIndex = 0; $fileIndex < $nGapDataFiles; $fileIndex++) {
	my $file = $gapDataFiles[$fileIndex];
	my $qualOffset = $gapDataQuals[$fileIndex];
	unless (exists($fileHandles{$file})) {
	    my $fh;
	    open ($fh,$file) || die "Couldn't open $file\n";
	    $fileHandles{$file} = $fh;
	}

	my $fh = $fileHandles{$file};
	while (my $line = <$fh>) {
	    chomp $line;
	    my ($scaff,$c1,$p1,$c2,$p2,$gapSize,$gapUnc,@reads) = split(/\t/,$line);

	    my ($contig1) = $c1 =~ /^(.+)\.[35]$/;
	    my ($contig2) = $c2 =~ /^(.+)\.[35]$/;
	    my $gapName = "$scaff.$contig1.$contig2";
	    unless (exists($gapIndex{$gapName})) {
		die "Error: unexpected gap $gapName\n";
	    }
	    my $gapNumber = $gapIndex{$gapName};

	    if ($goodIndex || ($gapNumber > $maxGapIndex)) {

		unless ($qualOffset == $defaultQualOffset) {
		    my @newReads = ();
		    foreach my $read (@reads) {
			my ($seq,$qual) = $read =~ /^([^:]+)\:(.+)$/;
			my @q = split(/ */,$qual); 
			my @Q = map { chr(ord($_) -  $qualOffset + $defaultQualOffset)} @q; 
			my $newQual = join("",@Q);
			push(@newReads, $seq . ":" . $newQual);
		    }
		    $line = join("\t",$scaff,$c1,$p1,$c2,$p2,$gapSize,$gapUnc,@newReads);
		    @reads = @newReads;
		}
	    }

	    if ($gapNumber < $minGapIndex) {
		die "Error: gap index ($gapNumber) less than minimum ($minGapIndex)\n";
	    } elsif ($gapNumber > $maxGapIndex) {
		push(@overflow,$line);
		last;
	    } else {
		if (exists($gapData{$gapNumber})) {
		    push(@{$gapData{$gapNumber}},@reads);
		} else {
		    $gapData{$gapNumber} = [$scaff,$c1,$p1,$c2,$p2,$gapSize,$gapUnc,@reads];
		}
	    }
	}
    }
    
    if ($goodIndex) {

	for (my $gapNumber = $minGapIndex; $gapNumber <= $maxGapIndex; $gapNumber++) {
	    if (exists($gapData{$gapNumber})) {
		my @gapInfo = @{$gapData{$gapNumber}};
		my $gapString = join("\t",@gapInfo);
		$nReadsDumped += scalar(@gapInfo) - 7;  # Seven non-read fields
		print OUT "$gapString\n";
	    }
	}

	close OUT;
	print STDERR "\tDone.  Found $nReadsDumped reads for gap closing.\n";
    }

    $gapBlockIndex++;

}

while (my ($file,$fh) = each(%fileHandles)) {
    close $fh;
}



