#!/usr/bin/perl -w

use strict;
use Getopt::Long;
use Cwd;

my $version = 0.1;

my $q = 20;
my $ref = "";
my $time = 0;
my $force = 0;
my $noNorm = 0;
my $normCov = 50;
my $species = "";
my $single = "";
my $reads1 = "";
my $reads2 = "";
my $baseDir = getcwd;
my $verFlag;
my $help = 0;
my $h = 0;
my $pe = 0;

GetOptions(	"assemblyName=s" => \$species,
			"ref=s" => \$ref,
			"q=i" => \$q,
			"time" => \$time,
			"force" => \$force,
			"noNorm" => \$noNorm,
			"normCov" => \$normCov,
			"single=s" => \$single,
			"reads1=s" => \$reads1,
			"reads2=s" => \$reads2,
			"outputDir=s" => \$baseDir,
			"help" => \$help,
			"h" => \$h,
			"version" => \$verFlag
);

if ($h or $help or not ($single or ($reads1 and $reads2))) {
	print "Please specify either a single merged file of single end reads (using --single) or two merged files for paired-end reads (using --reads1 and --read2)\n";
	print "\nOther options:\n--species\t\tSpecify the species being assembled to help organize output files\n";
	print "--q\t\t\tChange q-value threshold for reads filtering (default q20)\n";
	print "--baseDir\t\tChange base directory for output files (default is current working directory)\n";
	print "--ref\t\tPath to fasta for reference genome\n";
	print "--force\t\tIgnore the log file and always run entire pipeline\n";
	print "--noNorm\t\tSkip the digital normalization of filtered reads\n";
	print "--normCov\t\tChange the max kmer coverage for digital normalization. Lowering coverage speeds up assembly time, but increases risk of artifacts (default 50)\n";
	print "--time\t\tRecords the amount of time each stage of the pipeline takes\n";
	print "--version\t\tPrints pipeline version and dies\n";
	print "--h or --help\tPrints this message and quits\n\n";
	die;
}

if ($verFlag) {
	print "Pipeline version $version\n";
	die;
}

print "Checking read length\n";
my $readLength = 0;
if ($single) {
	print "Checking single end reads\n";
	$pe = 0;
	my $S;
	open ($S, $single) or die "Can't open single end reads at $single, please check the file location: $!\n";
	while (my $line = readline($S)) {
		next unless $line =~ /^@/;
		my $r = readline($S);
		chomp $r;
		$readLength = length($r);
		print "Detected read length of $readLength\n";
		last;
	}
	close $S;
}

if ($reads1 and $reads2) {
	print "Checking paired-end reads\n";
	$pe = 1;
	my ($P1, $P2);
	open ($P1, $reads1) or die "Can't open paired end reads 1 at $reads1, please check the file location: $!\n";
	open ($P2, $reads2) or die "Can't open paired end reads 2 at $reads2, please check the file location: $!\n";
	my ($rl1, $rl2);
	while (my $line = readline($P1)) {
		next unless $line =~ /^@/;
		my $r = readline($P1);
		chomp $r;
		$rl1 = length($r);
		print "Detected read length of $rl1 for reads1\n";
		last;
	}
	while (my $line = readline($P2)) {
		next unless $line =~ /^@/;
		my $r = readline($P2);
		chomp $r;
		$rl2 = length($r);
		print "Detected read length of $rl2 for reads2\n";
		last;
	}
	die "Can't determine read length, please check the input files" unless ($rl1 and $rl2);
	if ($rl1 != $rl2) {
		warn "Read lengths in $reads1 and $reads2 do not match.  Please check the files\n";
	}
	$readLength = $rl1;
	close $P1;
	close $P2;
}

my @q = split /,/, $q;

my $filReads = "$baseDir/FilteredReads/$species";
my $normReads = "$baseDir/NormalizedReads/$species";
my $asReads = "$baseDir/AssembledReads/$species";
my $runDir = getcwd;
my $log = "$asReads/log.txt";
my ($LOG, $TIME);
my %status;

print "Creating output directories\n";

if (-d "$baseDir") {
	print "Found $baseDir\n";
} else {
	mkdir ("$baseDir");
	print "Made $baseDir\n";
}

if (-d "$baseDir/AssembledReads/") {
	print "Found $baseDir/AssembledReads/\n";
} else {
	mkdir ("$baseDir/AssembledReads/");
	print "Made $baseDir/AssembledReads/\n";
}

if (-d "$baseDir/AssembledReads/$species/") {
	print "Found $baseDir/AssembledReads/$species/\n";
} else {
	mkdir ("$baseDir/AssembledReads/$species/");
	print "Made $baseDir/AssembledReads/$species/\n";
}

print "Opening log\n";
if (-e $log and $force != 1) {
	open ($LOG, "$log") or die "$log exists, but can't be read: $!\n";
	while (my $line = readline($LOG)) {
		chomp $line;
		next unless $line;
		
		my ($pr, $st) = split /\t/, $line;
		$status{$pr} = $st;
	}
	close ($LOG) or die "Can't close $log: $!\n";
	open ($LOG, ">>$log") or die "Can't append $log: $!\n";
} else {
	open ($LOG, ">$log") or die "Can't write to $log: $!\n";
}

my %finished;
print "Checking previously completed jobs\n";
foreach my $pr (sort {$a cmp $b} keys %status) {
	my @info = split /\-/, $pr;
	my $con;
	if ($info[2]) {
		$con = "$info[1]-$info[2]";
	} else {
		$con = $info[1];
	}
	
	if ($finished{$info[0]}) {
		$finished{$info[0]} .= ",$con";
	} else {
		$finished{$info[0]} = $con;
	}
}

if ($time) {
	open ($TIME, ">$asReads/time.txt");
}

my $efStartTime = time;
foreach my $q (@q) {
	if ($status{"erneFilter-q$q"}) {
		print "Previously finished erne-filter-q$q, skipping this run\n";
	} else {	
		if (-d "$filReads") {
			print "Found $filReads\n";
		} else {
			mkdir ("$filReads");
			print "Made $filReads\n";
		}
		if (-d "$filReads/q$q") {
			print "Found $filReads/q$q\n";
		} else {
			mkdir ("$filReads/q$q");
			print "Made $filReads/q$q\n";
		}
		if ($pe) {
			system("erne-filter --query1 $reads1 --query2 $reads2 --output-prefix $filReads/q$q/merged --ultra-sensitive --force-standard");
		} else {
			system("erne-filter --query1 $single --output-prefix $filReads/q$q/merged --ultra-sensitive --force-standard");
		}
		print $LOG "erneFilter-q$q\t1\n";
	}
}
my $efEndTime = time;
my $efRunTime = $efEndTime - $efStartTime;
print $TIME "ErneFilter:\t$efRunTime\n" if $time;

my $khStartTime = time;
my $finK = $finished{"khmer"};
$finK = "none" unless $finK;
foreach my $q (@q) {		
	if (-d "$normReads") {
		print "Found $normReads\n";
	} else {
		mkdir ("$normReads");
		print "Made $normReads\n";
	}
}

if ($finK =~ /q$q/) {
	print "Previously finished khmer, skipping this run\n";
} else {
	if ($pe) {
		print "Interleaving reads\n";
		system("interleave-reads.py  -o $filReads/q$q/mergedI.fastq --force $filReads/q$q/merged_1.fastq $filReads/q$q/merged_2.fastq");	
		print "Normalizing reads\n";
		if ($noNorm) {
			system("cp $filReads/q$q/mergedI.fastq $normReads/mergedI.fq");	
		} else {
			system("normalize-by-median.py -o $normReads/mergedI.fq -x 1000000000 -C $normCov --n_tables 99 --force --ksize 32 $filReads/q$q/mergedI.fastq");	
		}
		print "Splitting reads\n";
		chdir("$normReads/");
		system("extract-paired-reads.py $normReads/mergedI.fq");
		system("split-paired-reads.py --force $normReads/mergedI.fq.pe");
		system("mv $normReads/mergedI.fq.pe.1 $normReads/merged.fq.1");
		system("mv $normReads/mergedI.fq.pe.2 $normReads/merged.fq.2");
		system("mv $normReads/mergedI.fq.pe $normReads/mergedI.fq");
	} else {
		if ($noNorm) {
			system("cp $filReads/q$q/merged_1.fastq $normReads/merged.fq");
		} else {
			system("normalize-by-median.py -o $normReads/merged.fq -x 1000000000 -C 50 --n_tables 99 --force $filReads/q$q/merged_1.fastq");
		}
		system("fastq-to-fasta.py -o $normReads/merged.fa $normReads/merged.fq");
	}
	print $LOG "khmer-q$q\t1\n";
}
my $khEndTime = time;
my $khRunTime = $khEndTime - $khStartTime;
print $TIME "Khmer:\t$khRunTime\n" if $time;
chdir("$runDir/");

my $soStartTime = time;
#Prepare reference
system("bowtie2-build $ref $ref");
system("tophat2 -o $baseDir/AssembledReads/$species/ -p 8 $ref $normReads/merged.fq.1 $normReads/merged.fq.2");

my %nucl;
system("bayesembler -b $baseDir/AssembledReads/$species/accepted_hits.bam -o $baseDir/AssembledReads/$species/bayesembler-");
readAssembly("$baseDir/AssembledReads/$species/bayesembler-assembly", "bayesembler");
system("cufflinks -o $baseDir/AssembledReads/$species/cufflinks $baseDir/AssembledReads/$species/accepted_hits.bam");
readAssembly("$baseDir/AssembledReads/$species/cufflinks/transcripts", "cufflinks");
system("scallop -i $baseDir/AssembledReads/$species/accepted_hits.bam -o $baseDir/AssembledReads/$species/scallop.gtf");
readAssembly("$baseDir/AssembledReads/$species/scallop", "scallop");
system("stringtie -o $baseDir/AssembledReads/$species/stringtie.gtf $baseDir/AssembledReads/$species/accepted_hits.bam");
readAssembly("$baseDir/AssembledReads/$species/stringtie", "stringtie");

my ($OUT, $CONNUCL, $CONNUCLONG);
open ($OUT, ">$baseDir/AssembledReads/$species/consemble3+g.faa") or die "Can't open $baseDir/AssembledReads/$species/consemble3+g.faa: $!";
open ($CONNUCL, ">$baseDir/AssembledReads/$species/consemble3+g.fa") or die "Can't open $baseDir/AssembledReads/$species/consemble3+g.faa: $!";
open ($CONNUCLONG, ">$baseDir/AssembledReads/$species/consemble3+g_longest.fa") or die "Can't open $baseDir/AssembledReads/$species/consemble3+g.faa: $!";
my $i = 1;
my %seqNames;
my %seqCounts;
foreach my $seq (keys %seqCounts) {
	next unless $seq;
	if ($seqCounts{$seq} > 2) {
		my $shortest = 999999;
		my $longest = 0;
		my @assem = split /,/, $seqNames{$seq};
		my ($nuclSeq, $nuclSeqLong);
		foreach my $s (@assem) {
			#print "$s\n";
			if (length($nucl{$s}) < $shortest) {
				$nuclSeq = $nucl{$s};
				$shortest = length($nucl{$s});
			}
			if (length($nucl{$s}) > $longest) {
				$nuclSeqLong = $nucl{$s};
				$longest = length($nucl{$s});
			}
		}
		print $OUT ">Contig_$i\n$seq\n";
		print $CONNUCL ">Contig_".$i."\n$nuclSeq\n";
		print $CONNUCLONG ">Contig_".$i."\n$nuclSeqLong\n";
		$i++;
	}
}

sub readAssembly {
	my ($f, $a) = @_;
	system("gtf_to_fasta $f.gtf $ref $f.fa");
	readNucAssembly($f, $a);
	system("ORFfinder -in $f.fa -out $f-ORFs.fa");
	getLongestORF($f);
	my $F;
	open ($F, "$f-ORFs.fa.faa") or die "Can't open $f-ORFs.fa.faa: $!\n";
	my %seen;
	my $seq = "";
	my $id;
	while (my $line = readline($F)) {
		chomp $line;
		next unless $line;
		
		if ($line =~ /^>/) {
			if (!$seen{$seq}) {
				$seqCounts{$seq}++;
				$seen{$seq};
			}
			$line =~ s/^>//;
			my @info = split /\s+/, $line;
			if ($seqNames{$seq}) {
				$seqNames{$seq} .= ",$id" if $seq;
			} else {
				$seqNames{$seq} = "$id" if $seq;
			}
			$seq = "";
			$info[0] =~ s/\_\d+//;
			$id = "$a-$info[0]";
			print "$id\n";
		} else {
			$seq .= $line;
		}
	}
	if (!$seen{$seq}) {
		$seqCounts{$seq}++;
		$seen{$seq};
	}
	if ($seqNames{$seq}) {
		$seqNames{$seq} .= ",$id";
	} else {
		$seqNames{$seq} = "$id";
	}
	$seq = "";
}
sub readNucAssembly {
        my ($f, $a) = @_;

        my $seq = "";
        my $id;
		my $F;
		open ($F, "$f.fa") or die "Can't open $f.fa: $!\n";
        while (my $line = readline($F)) {
                chomp $line;
                next unless $line;

                if ($line =~ /^>/) {
                    $line =~ s/^>//;
                    my @info = split /\s+/, $line;
                    $nucl{$id} = $seq if $seq;
                    $seq = "";
                    $id = "$a-$info[0]";
					print "$id\n";
                } else {
                    $seq .= $line;
                }
        }
        $nucl{$id} = $seq if $seq;
        $seq = "";
}

sub getLongestORF {
	my $f = shift @_;
	my %tmpGenes;
	my ($tmpId, $tmpSeq);
	my ($TMP, $TMPOUT);
	open ($TMP, "$f-ORFs.fa") or die "Can't open$f-ORFs.fa\n";
	open ($TMPOUT, ">$f-ORFs.fa.faa") or die "Can't open $f-ORFs.fa.faa\n";

	while (my $line = readline($TMP)) {
		chomp $line;
		next unless $line;
		if ($line =~ /^>/) {
			if ($tmpSeq) {
				$tmpGenes{$tmpId} = "" unless $tmpGenes{$tmpId};
				$tmpGenes{$tmpId} = $tmpSeq if (length($tmpSeq) > length($tmpGenes{$tmpId}));
			}
			my @info = split /ORF\d+_/, $line;
			my @gn = split /:/, $info[1];
			$tmpId = $gn[0];
			$tmpSeq = "";
		} else {
			$tmpSeq .= $line;
		}
	}
	$tmpGenes{$tmpId} = $tmpSeq if (length($tmpSeq) > length($tmpGenes{$tmpId}));

	foreach my $g (sort {$a cmp $b} keys %tmpGenes) {
		print $TMPOUT ">$g\n$tmpGenes{$g}\n";
	}

	close ($TMP);
	close ($TMPOUT);
}