#!/usr/bin/perl
BEGIN{ push(@INC, $ENV{LOCAL_PERL_MODULES}); }
use strict;
use warnings;
use Getopt::Std;
use FASTA::FileIO;
use FASTA::Util;
use BLAST::Util;
use BLAST::M9FileIO;
use Sequence::LociUtil;
use MetaGene::FileIO;
use MetaGene::Util;
use MetaGene::Status;
use Utils::Calc;
use vars qw($opt_h $opt_i $opt_n $opt_d $opt_p $opt_L $opt_E $opt_C $opt_P);
&getopts('hi:n:d:p:L:E:P:');

my $usage = <<_EOH_;

## Options ###########################################
## Required:
# -i    a scaffold FASTA file name
# -n    job name
# -d    blast database file (dna)
# -p    blast database file (pep)

## Optional:
# -L    minimum scaffolds/contigs length (default = 100)
# -E    maximum E-value (default = 1e-5)
# -P    num. threads (default = 1)
# -C    chimera rate cutoff (default = 1%)

## Others:
# -h print help

_EOH_
    ;

die $usage if $opt_h;

# Get command line options
my $minLength   = 100;  $minLength   = $opt_L if defined $opt_L;
my $maxEval     = 1e-5; $maxEval     = $opt_E if defined $opt_E;
my $numCPUs     = 1;    $numCPUs     = $opt_P if defined $opt_P;
my $chimeraRate = 0.01; $chimeraRate = $opt_C if defined $opt_C;
my $inFile      = $opt_i or die $usage;
my $jobName     = $opt_n or die $usage;
my $dnaDB       = $opt_d or die $usage;
my $pepDB       = $opt_p or die $usage;

# Print parameters
printParameters();

# Get scaffold FastA sequences ...
my $scafFile = $jobName . '.scafs.fa';
my $scafSeqs = [];
if( -e $scafFile ){
    $scafSeqs = FASTA::FileIO::parse( $scafFile );
    print STDERR scalar(@$scafSeqs) . " scaffold sequences were loaded.\n";
} else {
    $scafSeqs = FASTA::FileIO::parse( $inFile );
    foreach my $seq (@$scafSeqs){
	my $title = $seq->getTitle();
	$title    =~ s/\s+/_/g;
	$seq->setTitle( $title );
    }
    print STDERR scalar(@$scafSeqs) . " scaffold sequences were imported.\n";
    FASTA::FileIO::output( $scafFile, $scafSeqs );
}

# Get long scaffold FastaA sequences ...
my $longScafFile = $jobName . '.long-scafs.fa';
my $longScafSeqs = [];
if( -e $longScafFile ){
    $longScafSeqs = FASTA::FileIO::parse( $longScafFile );
    print STDERR scalar(@$scafSeqs) . " long scaffold sequences were loaded.\n";
} else {
    $longScafSeqs = FASTA::Util::extractLongSequences( $scafSeqs, $minLength );
    print STDERR scalar(@$longScafSeqs) . " long scafffold sequences were extracted.\n";
    FASTA::FileIO::output( $longScafFile, $longScafSeqs );
}

# Get contig FastA sequences ...
my $contigFile = $jobName . '.contigs.fa';
my $contigSeqs = [];
if( -e $contigFile ){
    $contigSeqs = FASTA::FileIO::parse( $contigFile );
    print STDERR scalar(@$contigSeqs) . " contig sequences were loaded.\n";
} else {
    $contigSeqs = FASTA::Util::scaffolds2contigs( $scafSeqs );
    print STDERR scalar(@$contigSeqs) . " contig sequences were obtained.\n";
    FASTA::FileIO::output( $contigFile, $contigSeqs );
}

# Get long contig FastA sequences ...
my $longContigFile = $jobName . '.long-contigs.fa';
my $longContigSeqs = [];
if( -e $longContigFile ){
    $longContigSeqs = FASTA::FileIO::parse( $longContigFile );
    print STDERR scalar(@$longContigSeqs) . " long contig sequences were loaded.\n";
} else {
    $longContigSeqs = FASTA::Util::extractLongSequences( $contigSeqs, $minLength );
    print STDERR scalar(@$longContigSeqs) . " long contig sequences were extracted.\n";
    FASTA::FileIO::output( $longContigFile, $longContigSeqs );
}

# Compute contig stats
my $contigStats = calcStats( $contigSeqs, $minLength );
print "Contig stats...\n";
printStats( $contigStats );

# Compute scaffold stats
my $scafStats = calcStats( $scafSeqs, $minLength );
print "Scaffold stats...\n";
printStats( $scafStats );

# Compute contig cover and chimera rates
my $contigRates = calcRates( $jobName.'.long-contigs', $dnaDB, $longContigFile, 
			     $maxEval, $numCPUs, $minLength, $chimeraRate );
print " Contig rates...\n";
printRates( $contigRates );

# Compute scaf cover and chimera rates
my $scafRates = calcRates( $jobName.'.long-scafs', $dnaDB, $longScafFile, 
			   $maxEval, $numCPUs, $minLength, $chimeraRate );
print " Scaffold rates...\n";
printRates( $scafRates );

=comment
# Find genes from contigs
my $contigGeneStats = calcGeneStats( $jobName.'.long-contigs', $pepDB, 
				     $longContigFile );
print " Contig gene stats...\n";
printGeneStats( $contigGeneStats );

# Find genes from scaffolds
my $scafGeneStats = calcGeneStats( $jobName.'.long-scafs', $pepDB, 
				   $longScafFile );
print " Contig gene stats...\n";
printGeneStats( $scafGeneStats );
=cut


sub printParameters {
    print "\n"
	. "*** parameters ***************************************" . "\n"
	. " Scaffold flie           = " . $inFile    . "\n"
	. " Minimum sequence length = " . $minLength . "\n"
	. "******************************************************" . "\n\n";
}

sub calcStats {
    my $seqs = shift;
    my $minLength = shift;
    my $longSeqs = FASTA::Util::extractLongSequences( $seqs, $minLength );
    my $stats = {};
    $stats->{numSeqs} = scalar( @$seqs );
    $stats->{numLongSeqs} = scalar( @$longSeqs );
    $stats->{totalLength} = FASTA::Util::calcTotalLength( $longSeqs );
    $stats->{n50} = FASTA::Util::calcNStats( $longSeqs, 50 );
    $stats->{n90} = FASTA::Util::calcNStats( $longSeqs, 90 );
    return $stats;
}

sub printStats {
    my $stats = shift;
    print "\n"
	. "*** assembly stats ***********************************" . "\n"
	. " Num. sequencess        = " . $stats->{numSeqs} . "\n"
	. " Num. long sequences    = " . $stats->{numLongSeqs} . "\n"
	. " Total sequence length  = " . $stats->{totalLength} . "\n"
	. " N50 size               = " . $stats->{n50} . "\n"
	. " N90 size               = " . $stats->{n90} . "\n"
	. "******************************************************" . "\n\n";
}

sub calcRates {
    my $jobName     = shift;
    my $dnaDB       = shift;
    my $queryFile   = shift;
    my $maxEval     = shift;
    my $numCPUs     = shift;
    my $minLength   = shift;
    my $chimeraRate = shift;
    my $blsFile     = $jobName . '.blast';
    unless( -e $blsFile ){
	print STDERR "Running blast (vs. dnaDB) ...\n";
	BLAST::Util::runBlast( { programName => 'blastn',
				 formatCode  => 8,
				 queryFile   => $queryFile, 
				 subjectFile => $dnaDB, 
				 outputFile  => $blsFile, 
				 maxEvalue   => $maxEval, 
				 numCPUs     => $numCPUs
			       } );
    }
    my $rates = {};
    $rates->{cover} = calcCoverRate( $blsFile, $dnaDB, 
				     $jobName.'.cover-rate', $minLength );
    $rates->{chimera} = calcChimeraRate( $blsFile, $queryFile, 
					 $jobName.'.chimera-rate', $chimeraRate );
    return $rates;
}

sub calcCoverRate {
    my $blsFile = shift;
    my $dbFile  = shift;
    my $outFile = shift;
    my $minLength = shift;
    print STDERR "Calculating cover rate ...\n";
    my $seqs = FASTA::FileIO::parse($dbFile);
    print STDERR scalar(@$seqs) . " database sequences were loaded.\n";
    my $mergedLoci = getCoveredLoci( $blsFile, $minLength );

    # Compute cover rates
    open OUT, ">$outFile" or die "Can't open $outFile to write.";
    my $totalRefLength = 0;
    my $totalCoverLength = 0;
    foreach my $seq (@$seqs) {
	my $title = $seq->getShortTitle();
	my $refLength = $seq->getLength();
	my $curLoci = Sequence::LociUtil::selectBySeqID($mergedLoci, $title);
	my $coverLength = Sequence::LociUtil::getTotalLength($curLoci);
	my $coverRate = Utils::Calc::round( $coverLength/$refLength*100, 2 );
	print OUT "\n"
	    . "--------------------------------------------------------------\n"
	    . $seq->getTitle() . "\n"
	    . $coverLength . " / " . $refLength 
	    . " ( " . $coverRate . " %)" . "\n"
	    . "--------------------------------------------------------------\n";
	$totalRefLength += $refLength;
	$totalCoverLength += $coverLength;
    }
    
    my $coverRate = Utils::Calc::round( $totalCoverLength/$totalRefLength*100, 2 );
    print OUT "\n"
	. "--------------------------------------------------------------\n"
	. "Total cover rate" . "\n"
	. $totalCoverLength . " / " . $totalRefLength
	. " ( " . $coverRate . " %)" . "\n"
	. "--------------------------------------------------------------\n";
    close OUT;
    print STDERR "...done.\n\n";
    return $coverRate;
}

sub getCoveredLoci {
    my $blsFile = shift;
    my $minLength = shift;
    my $hits = BLAST::M9FileIO::parse($blsFile);
    print STDERR scalar(@$hits) . " BLAST hits were loaded.\n";
    my $longHits = BLAST::Util::extractLongHits($hits, $minLength);
    print STDERR scalar(@$longHits) . " long hits were found.\n"
	. "\tlength cutoff value = $minLength.\n";
    my $refLoci = BLAST::Util::getReferenceLoci($longHits);
    print STDERR scalar(@$refLoci) . " reference loci were covered.\n";
    my $mergedLoci = Sequence::LociUtil::getMergedLoci($refLoci);
    print STDERR scalar(@$mergedLoci) . " unoverlapped reference loci were covered.\n";
    return $mergedLoci;
}

sub calcChimeraRate {
    my $blsFile = shift;
    my $faFile  = shift;
    my $outFile = shift;
    my $cRate   = shift;
    print STDERR "Calculating chimera rate...\n";
    # Load BLAST hits ...
    my $hits = BLAST::M9FileIO::parse($blsFile);
    print STDERR scalar(@$hits) . " BLAST hits were loaded.\n";

    # Load scaffold sequences ...
    my $seqs = FASTA::FileIO::parse($faFile);
    print STDERR scalar(@$seqs) . " queried sequences were loaded.\n";
    my $totalScaffoldLength = FASTA::Util::calcTotalLength($seqs);

    open OUT, ">$outFile" or die "Can't open $outFile to write.";
    # Compute chimeric scaffolds length
    my $queryID2hits = getQueryID2hits($hits);
    my $queryID2length = getQueryID2length($seqs);
    print STDERR scalar(keys %$queryID2hits) . " query IDs were found.\n";
    my $numChimericScaffolds = 0;
    my $sumChimericHitLength = 0;
    my $sumChimericScaffoldLength = 0;
    my $pair2length = {};
    foreach my $queryID (sort keys %$queryID2hits){
	my $queryLength = $queryID2length->{$queryID};
	my $curHits = $queryID2hits->{$queryID};
	my $bestHit = getBestHit($curHits);
	my $chimericHits = getChimericHits($curHits,$bestHit,$queryLength*$cRate);
	if( scalar(@$chimericHits) ){
	    ++$numChimericScaffolds;
	    my $chimeraLoci = BLAST::Util::getQueryLoci($chimericHits);
        my $mergedChimeraLoci =Sequence::LociUtil::getMergedLoci($chimeraLoci);
        $sumChimericHitLength +=Sequence::LociUtil::getTotalLength($mergedChimeraLoci);   
	    #$sumChimericHitLength += getTotalQueryLength($chimericHits);
	    $sumChimericScaffoldLength += $queryLength;
	    print OUT "Chimera: " . $queryID 
		. "\t" . $bestHit->getSubjectID()
		. "\t" . $sumChimericHitLength . "\n";
	}
	foreach my $chimericHit (@$chimericHits){
	    my $pair = $bestHit->getSubjectID() . '-' . $chimericHit->getSubjectID();
	    if( $pair2length->{$pair} ){
		$pair2length->{$pair} += $chimericHit->getQueryLocus()->getLength();
	    } else {
		$pair2length->{$pair} = $chimericHit->getQueryLocus()->getLength();
	    }
	}
    }

    foreach my $pair (reverse sort { $pair2length->{$a} <=> $pair2length->{$b} } keys %$pair2length ){
	print OUT $pair . "\t" . $pair2length->{$pair} . "\n";
    }
    
    print OUT "\n"
	. "#. scaffolds                   = " . scalar(@$seqs) . "\n"
	. "#. chimeric scaffolds          = " . $numChimericScaffolds 
	. " ( " 
	. Utils::Calc::round($numChimericScaffolds/scalar(@$seqs)*100,2) 
	. " %)\n"
	. "\n"
	. "Total scaffold length          = " . $totalScaffoldLength . "\n"
	. "Total chimeric region length   = " . $sumChimericHitLength 
	. " ( " 
	. Utils::Calc::round($sumChimericHitLength/$totalScaffoldLength*100,2) 
	. " %)\n"
	. "Total chimeric scaffold length = " . $sumChimericScaffoldLength 
	. " ( " 
	. Utils::Calc::round($sumChimericScaffoldLength/$totalScaffoldLength*100,2) 
	. " %)"
	. "\n\n";
    close OUT;
    print STDERR "...done.\n\n";
    return Utils::Calc::round($sumChimericHitLength/$totalScaffoldLength*100,2);
}

sub getQueryID2hits {
    my $hits = shift;
    my $queryID2hits = {};
    foreach my $hit (@$hits){
	my $queryID = $hit->getQueryID();
	if( $queryID2hits->{$queryID} ){
	    push( @{$queryID2hits->{$queryID}}, $hit);
	} else {
	    $queryID2hits->{$queryID} = [$hit];
	}
    }
    return $queryID2hits;
}


sub getQueryID2length {
    my $seqs = shift;
    my $queryID2length = {};
    foreach my $seq (@$seqs){
	$queryID2length->{$seq->getShortTitle()} = $seq->getLength();
    }
    return $queryID2length;
}

sub getBestHit {
    my $hits = shift;
    my $bestHit = $hits->[0];
    foreach my $hit (@$hits){
	if( $hit->getBitScore() > $bestHit->getBitScore() ){
	    $bestHit = $hit;
	}
    }
    return $bestHit;
}


sub getChimericHits {
    my $hits = shift;
    my $bestHit = shift;
    my $cutLen = shift;
    my $chimericHits = [];
    my $nonChimericHits 
	= selectHitsBySubjectID($hits, $bestHit->getSubjectID());
    foreach my $hit (@$hits){
	if( $hit->getSubjectID() ne $bestHit->getSubjectID() ){
	    my $curQueryLocus = $hit->getQueryLocus();
	    my $flgChimeric = 1;
	    foreach my $nonChimericHit (@$nonChimericHits){
		my $nccQueryLocus = $nonChimericHit->getQueryLocus();
		if( $curQueryLocus->isOverlap($nccQueryLocus) ){
		    $flgChimeric = 0;
		    last;
		}
	    }
	    if( $flgChimeric ){
		push( @$chimericHits, $hit );
	    }
	}
    }
    return $chimericHits;
}


sub selectHitsBySubjectID {
    my $hits = shift;
    my $id = shift;
    my $selectedHits = [];
    foreach my $hit (@$hits){
	if( $hit->getSubjectID() eq $id ){
	    push( @$selectedHits, $hit );
	}
    }
    return $selectedHits;
}


sub getTotalQueryLength {
    my $hits = shift;
    my $totalLength = 0;
    foreach my $hit (@$hits){
	$totalLength += $hit->getQueryLocus()->getLength();
    }
    return $totalLength;
}

sub printRates {
    my $rates = shift;
    print "\n"
	. "*** cover and chimera rate ***************************" . "\n"
	. " Cover rate              = " . $rates->{cover}   . "\n"
	. " Chimera rate            = " . $rates->{chimera} . "\n"
	. "******************************************************" . "\n\n";
}

sub calcGeneStats {
    my $jobName = shift;
    my $pepDB   = shift;
    my $seqFile = shift;
    my $mgFile  = $jobName . '.MetaGene.raw';
    unless( -e $mgFile ){
	system "metagene $seqFile > $mgFile";
    }
    my $pepFile = $jobName . '.pep.fa';
    my $orfs = MetaGene::FileIO::parse( $mgFile );
    print STDERR scalar(@$orfs) . " ORFs were loaded.\n";
    my $pepSeqs = [];
    if( -e $pepFile ){
	$pepSeqs = FASTA::FileIO::parse( $pepFile );
	print STDERR scalar(@$pepSeqs) . " protein sequences were loaded.\n";
    } else {
	my $dnaSeqs = FASTA::FileIO::parse( $seqFile );
	print STDERR $dnaSeqs->[0]->getShortTitle() . "\n";
	print STDERR scalar(@$dnaSeqs) . " contig/scaffold sequences were loaded.\n";
	$pepSeqs = MetaGene::Util::translateORFs( $orfs, $dnaSeqs );
	print STDERR scalar(@$pepSeqs) . " protein sequences were output.\n";
	FASTA::FileIO::output( $pepFile, $pepSeqs );
    }
    return MetaGene::Status::calcStats( $orfs );
}

sub printGeneStats {
    my $stats = shift;
    print "\n"
	. "*** cover and chimera rate ***************************" . "\n";
    MetaGene::Status::printStats($stats);
    print "******************************************************" . "\n\n";
}
