#!/usr/bin/perl -w
# =======================================================================================
# secondary_assembly_stats.pl
# A script to compute depth of coverage and generate a list of reads for each contig for 
# a CAP3 assembly used to join primary contigs and singletons from MIRA. Output files can
# be used to generate plots and histograms of coverage or to extract all of the reads 
# used in the assembly.
#
# Joshua Der
#
# Version history
# =======================================================================================
# Version 1.3 - 19 March 2010
# -Summary output will be written to a file with a standardized name if explicitly 
#  requested on $ARGV[4] with the string "yes".
# -Detailed summary statistics for sequence length are generated identical to those 
#  produced by get_seq_len_stats.pl (dePamphilis lab @ PennState). 
#  Requires the Statistics::Descriptive perl module.
#
# Version 1.2 - 18 March 2010
# -Name changed to secondary_assembly_stats.pl to reflect new functionality.
# -Updated pattern matching to accomodate Sequence IDs extracted and modified by NCBI 
#  fastacmd or blastdbcmd (e.g. "lcl|contig_name). This enables you to use this script on
#  CAP3 assemblies using only a subset of the mira contigs.
# -Now requires the MIRA contig read list file as $ARGV[3] and produces a new version of 
#  this file for the subset of contigs used in the CAP3 assembly.
# 
#
# Version 1.1 - 11 February 2010
# -Added read count and mean unigene length to summary output.
#
# Version 1.0 - 28 January 2010
# -First release of coverage_secondary.pl
# =======================================================================================
use strict;

if (!$ARGV[3]) {
    print "USAGE: secondary_assembly_stats.pl <cap3.contigs.ace> <cap3.singlets> \\\n";
    print "       <mira_info_contigstats.txt> <mira_info_contigreadlist.txt> \"yes\"\n";
    print "       Summary output will be written to a file if the last argument is \"yes\".\n";
    exit(1);
}

# get the command line arguments
my $cap3_ace_file = $ARGV[0];
my $cap3_singlets = $ARGV[1];
my $mira_contig_stats = $ARGV[2];
my $mira_contig_read_list = $ARGV[3];
# name the output files
my $base_name = substr($cap3_ace_file, 0, - 4);
my $coverage_output = $base_name . '.all.coverage_stats.txt';
my $read_list_output = $base_name . '.all.read_list.txt';
my $len_stats = $base_name . '.all.seq_stats.txt';
# if we request the summary file, set it 
my $summary_file = undef;
if ($ARGV[4]) {
	if ($ARGV[4] eq "yes") {
		$summary_file = $base_name . '.all.summary.txt';
	}
}
# intialize global variables
my %cap3_hash = ();
my @singleton_list = ();
my @primary_contig_retained_list = ();
my %mira_stats = ();
my %mira_reads = ();
my @contig_coverage = ();
my @contig_length = ();
my $sum_est = 0;
my $total_read_length = 0;
my $total_contig_length = 0;
my $total_coverage = 0;
my $summary = "";

# open cap3 ACE file and build the cap3_hash
open(ACE, "$cap3_ace_file") || die("Cannot open $cap3_ace_file file");
# read ACE file
while ($_ = <ACE>){
    # skip until the line describes a contig
    if ($_ !~ /^CO\s+(\w+)\s+(\d+)\s+(\d+).*/){next;}
	# store the contig id
    my $cap3_contig_id = $1;
    # create a hash of contigs (key = contig_id) and put a hash of contig stats in it
    $cap3_hash{$cap3_contig_id} = {
    	'cap3_contig_length' => $2,
		'num_mira_contigs' => $3
		};
	# run through the contig sequence and count the pads
	my $pads = 0;
	while ($_ = <ACE>) {
		$pads += ($_ =~ tr/*//);
		if ($_ =~ /^BQ/) {last;}
	}
	# adjust the contig length for pads
	$cap3_hash{$cap3_contig_id}{'cap3_contig_length'} -= $pads;
	# adding a loop structure to collect the "read" names
    my @mira_contig = ();
    while ($_ = <ACE>) {
    	# skip until the line describes a read
		if ($_ !~ /^AF\s+(\S+)/){next;}
    	my $id =$1;
    	# check if the contig name has "lcl|" at the beginning and remove it
		if ($id =~ /^lcl\|(\S+)/) {
			$id = $1;
		}
		# store read names as an array that will be put in the hash of contig stats
		push (@mira_contig, $id);
		# exit this while loop once we have collected all of the primary contig ids
		if(scalar @mira_contig == $cap3_hash{$cap3_contig_id}{'num_mira_contigs'}) {last;};
	}
	# put the array of primary contig ids into the cap3 contig stats hash, note the need to use an array reference
	@{ $cap3_hash{$cap3_contig_id}{'mira_contig_ids'} } = @mira_contig;
}
close ACE;

# open cap3 singlets file (fasta format) and build an array of ids
open(SINGLET, "$cap3_singlets") || die("Cannot open $$cap3_singlets file");
while($_ = <SINGLET>){
	chomp $_;
	# look for the deflines
	if ($_=~/^>(\S+)/){
		my $id = $1;
		# check if the contig name has "lcl|" at the beginning and remove it
		if ($id =~ /^lcl\|(\S+)/) {
			$id = $1; #print "$id\n";
		}
		# if the name contains '_s' its a "true" singleton; this distinction allows us to count them at the end
		if ($id=~/(\S+_s\S+)/) {
			# collect all of the singletons in an array
			push @singleton_list, $id;
		}
		# if its not a "true" singleton:
		else {
			push @primary_contig_retained_list, $id;
		}
	}
}
close SINGLET;

# open MIRA contig stats file and build the mira_stats hash
open(MIRA_STATS, "$mira_contig_stats") || die("Cannot open $mira_contig_stats");
while (my $_ = <MIRA_STATS>){
	# skip the first line (it begins with '#')
	if ($_ =~ /^#/){next;}
	# split each line by whitespace and put it into an array
	my @mira_contig_stats = split(/\s+/, $_);
	# remove the first element of the array (contig_id) and put it in a new variable
	my $mira_contig = shift @mira_contig_stats;
	# build the mira hash with mira contig ids as the key and the stats array as the value (using an array reference)
	# order of values in the array is: length	av.qual	#-reads	mx.cov.	av.cov	GC%	CnIUPAC	CnFunny	CnN	CnX	CnGap	CnNoCov
	@{ $mira_stats{$mira_contig} } = @mira_contig_stats;
}
close MIRA_STATS;

# open coverage output file
open (COVERAGE, ">$coverage_output") || die ("Cannot open $coverage_output output file");
# This draws data from the cap3 and mira_stats hashes to print the output coverage file: unigene_id	length	#EST	x-fold_coverage
# It would be easy to add more fields to the output file if desired by modifying the lines begining with 'print COVERAGE'
print COVERAGE join("\t", '#Contig_ID', 'Contig_length', 'Num_EST', 'Sequence_coverage'), "\n";
# run through each cap3 contig and collect data from the cap3_hash and mira_hash
# for CAP3 contig ids named "Contig###", we sort by just the number :-) 
foreach my $cap3_contig (sort {(substr $a,6) <=> (substr $b,6)} keys %cap3_hash) {
	my $cap3_length = $cap3_hash{$cap3_contig}{'cap3_contig_length'}; 
	my $cap3_est = 0; # to calculate I sum the read counts from each mira contig in the cap3 contig
	my $sum_mira_read_lengths = 0;  # Equal to the sum of the product (mira_length * mira_coverage) from each mira contig. 
									# This is necessarily an estimate (based on the average coverage) without parsing the
									# RD tags in the ACE file and determining the unpadded length
	my $cap3_coverage = 0;  # to calculate I divide 'sum_mira_read_lengths' from each mira contig in
							# the cap3 contig and divide by the cap3 contig length 
	# run through each mira contig assembled into the cap3 contig
	foreach my $mira_contig (@{$cap3_hash{$cap3_contig}{'mira_contig_ids'}}) {
		# retreive the number of reads for the mira contig and add it to the total for the cap3 contig
		$cap3_est += @{$mira_stats{$mira_contig}}[2];
		# retreive the length of the mira contig
		my $mira_contig_length = @{$mira_stats{$mira_contig}}[0];
		# retreive the coverage of the mira contig
		my $mira_coverage = @{$mira_stats{$mira_contig}}[4];
		# calculate the total length of reads assembled into the cap3 contig
		$sum_mira_read_lengths += ($mira_contig_length * $mira_coverage);
	}
	# compute the contig coverage statistics
	$cap3_coverage = ($sum_mira_read_lengths / $cap3_length);
	# add the cap3 contig coverage to the unigene coverage array
	push @contig_coverage, $cap3_coverage;
	# add the cap3 contig length to the unigene length array
	push @contig_length, $cap3_length;
	# print stats to the output file
	print COVERAGE join("\t", $cap3_contig, $cap3_length, $cap3_est, sprintf ("%.2f", $cap3_coverage)), "\n";
	# increment the global coverage statistics
	$sum_est += $cap3_est;
	$total_read_length += $sum_mira_read_lengths;
	$total_contig_length += $cap3_length;
}
# collect stats for the primary contigs retained as unigenes
foreach my $mira_contig (@primary_contig_retained_list) {
	# get the number of ESTs for the contig
	my $mira_est += @{$mira_stats{$mira_contig}}[2];
	# retreive the length of the mira contig
	my $mira_contig_length = @{$mira_stats{$mira_contig}}[0];
	# retreive the coverage of the mira contig
	my $mira_coverage = @{$mira_stats{$mira_contig}}[4];
	# add it to the unigene coverage array
	push @contig_coverage, $mira_coverage;
	# add the contig length to the unigene length array
	push @contig_length, $mira_contig_length;
	# print stats to the output file
	print COVERAGE join("\t", $mira_contig, $mira_contig_length, $mira_est, $mira_coverage), "\n";
	# calculate read lengths for the contig
	my $mira_read_length = ($mira_contig_length * $mira_coverage);
	# increment the global coverage stats
	$sum_est += $mira_est;
	$total_read_length += $mira_read_length;
	$total_contig_length += $mira_contig_length;
}
# collect stats for the true singletons
foreach my $mira_singleton (@singleton_list) {
	# get the number of ESTs for the contig (=1)
	my $mira_est += @{$mira_stats{$mira_singleton}}[2];
	# retreive the length of the mira contig
	my $mira_singleton_length = @{$mira_stats{$mira_singleton}}[0];
	# retreive the coverage of the mira contig (=1)
	my $mira_coverage = @{$mira_stats{$mira_singleton}}[4];
	# add it to the unigene coverage array
	push @contig_coverage, $mira_coverage;
	# add the singleton length to the unigene length array
	push @contig_length, $mira_singleton_length;
	# print stats to the output file
	print COVERAGE join("\t", $mira_singleton, $mira_singleton_length, $mira_est, $mira_coverage), "\n";
	# increment the global coverage stats
	$sum_est += $mira_est;
	$total_read_length += $mira_singleton_length;
	$total_contig_length += $mira_singleton_length;
}
close COVERAGE;


# open MIRA contig read list file and build the mira_reads hash
open(MIRA_READS, "$mira_contig_read_list") || die("Cannot open $mira_contig_read_list");
while (<MIRA_READS>){
	# skip the comment lines (it begins with '#')
	if ($_ =~ /^#/){next;}
	chomp;
	# split each line by whitespace and put it into an array
	my @mira_contig_reads = split(/\s/, $_);
	my ($mira_contig, $read) = @mira_contig_reads;
	# build a mira hash with mira contig ids as the key and an array of reads as the value
	push @{ $mira_reads{$mira_contig} }, $read;
}
close MIRA_READS;

# open read list output file
open (READS_OUT, ">$read_list_output") || die ("Cannot open $read_list_output output file");
# print the file headers
print READS_OUT join("\t", '#Contig_ID', 'Read_name'), "\n";
# run through each cap3 contig and print all of the reads to file
foreach my $cap3_contig (sort {(substr $a,6) <=> (substr $b,6)} keys %cap3_hash) {
	# run through each mira contig assembled into the cap3 contig
	foreach my $mira_contig ( @{ $cap3_hash{$cap3_contig}{'mira_contig_ids'} } ) { #print "$cap3_contig\t$mira_contig\n";
		foreach my $mira_read ( @{ $mira_reads{$mira_contig} } ) {
			print READS_OUT "$cap3_contig\t$mira_read\n";
		}
	}
}
# print reads for the primary contigs retained in the CAP3 assembly
foreach my $mira_contig (@primary_contig_retained_list) {
	foreach my $mira_read ( @{ $mira_reads{$mira_contig} } ) {
		print READS_OUT "$mira_contig\t$mira_read\n";
	}
}
# and now handle the true singletons
foreach my $mira_singleton (@singleton_list) {
	# print stats to the output file
	foreach my $mira_read ( @{ $mira_reads{$mira_singleton} } ) { # a bit redundant, but safe in case the contig naming scheme in mira changes
		print READS_OUT "$mira_singleton\t$mira_read\n";
	}
}
close READS_OUT;

# Calculate some assembly summary statistics

# count the number of final sequences
my $num_unigenes = scalar @contig_coverage;
# count the number of primary contigs retained
my $num_primary_contigs = scalar @primary_contig_retained_list;
# count the number of singletons
my $num_singletons = scalar @singleton_list;
# count the number of secondary contigs
my $num_secondary_contigs = $num_unigenes - $num_primary_contigs - $num_singletons;

# calculate the unweighted average coverage for the final sequence list (does not consider contig length); probably a less informative statistic
my $sum_coverage = 0;
foreach (@contig_coverage) { $sum_coverage += $_; }
my $average_contig_coverage = sprintf ("%.2f", ($sum_coverage / $num_unigenes) );

# calculate the mean length for the final sequence list
my $sum_length = 0;
foreach (@contig_length) { $sum_length += $_; }
my $average_contig_length = sprintf ("%.2f", ($sum_length / $num_unigenes) );

# calculate the weighted average coverage for the unigene set (considers contig length)
$total_coverage = $total_read_length / $total_contig_length;
$total_coverage = sprintf ("%.2f", $total_coverage);

# Generate a descriptive stats for sequence length
# code taken from get_seq_len_stats.pl
use Statistics::Descriptive;
my $stat = Statistics::Descriptive::Full->new();
$stat->add_data(@contig_length);
my $seqcount = $stat->count();
my $mean = sprintf("%.2f", $stat->mean());
my $median = sprintf("%.2f", $stat->median());
my $variance = sprintf("%.2f", $stat->variance());
my $stdev = sprintf("%.2f", $stat->standard_deviation());
my $mode = sprintf("%.2f", $stat->mode());
my $min = sprintf("%.2f", $stat->min());
my $max = sprintf("%.2f", $stat->max());
my $sum = sprintf("%.0f", $stat->sum());
my $mb = sprintf("%.2f", $sum/1000000);
open LENGTHS_STATS, ">$len_stats";
print LENGTHS_STATS "assembly\tseqcount\tmean\tstdev\tmedian\tmode\tmin\tmax\tsum\tmb\n";
print LENGTHS_STATS "$base_name\t$seqcount\t$mean\t$stdev\t$median\t$mode\t$min\t$max\t$sum\t$mb\n";
close LENGTHS_STATS;

# Generate the coverage stats summary
$summary .= "\n\n****************************\n* Coverage summary results *\n****************************\n";
$summary .= "Input ACE file: \"$cap3_ace_file\"\n";
$summary .= "Input singlets file: \"$cap3_singlets\"\n";
$summary .= "Input contig stats file: \"$mira_contig_stats\"\n";
$summary .= "Input contig reads file: \"$mira_contig_read_list\"\n\n";
$summary .= "Total number of sequences: $num_unigenes\n";
$summary .= "Total number of reads assembled: $sum_est\n";
$summary .= "Total assembly length: $sum_length bp\n";
$summary .= "Mean sequence length: $average_contig_length\n";
$summary .= "Sequence subcategories:\n";
$summary .= "\tTotal number of secondary contigs: $num_secondary_contigs\n";
$summary .= "\tTotal number of primary contigs: $num_primary_contigs\n";
$summary .= "\tTotal number of singletons: $num_singletons\n\n";
$summary .= "Average sequence coverage: $average_contig_coverage" . "X\n";
$summary .= "\t(Weights each sequence equally, regardless of length.\n";
$summary .= "\tSingletons and a high frequency of low coverage contigs\n\tdrives this statistic down.)\n\n";
$summary .= "Total overall assembly coverage: $total_coverage" . "X\n";
$summary .= "\t(Equivalent to the average per-base read depth of coverage.)\n\n";
$summary .= "Output coverage stats file: \"$coverage_output\"\n";
$summary .= "\tOutput format: " . join(" ", '#Sequence_ID', 'Sequence_length', 'Num_EST', 'Sequence_coverage') . "\n\n";
$summary .= "Output read list file: \"$read_list_output\"\n";
$summary .= "\tOutput format: " . join("\t", '#Sequence_ID', 'Read_name') . "\n\n";

if ($summary_file) {
	open (SUMMARY, ">$summary_file")  || die ("Cannot open $summary_file output file");
	print SUMMARY $summary;
	close SUMMARY;
}
else {
	print $summary;
	print "Note: This results summary can be written to file if \"yes\" is given as \n";
	print "      the fifth argument when you call the script\n\n";
}
exit;