#!/usr/bin/env perl

# Copyright (c) 2017 Michael Roach (Australian Wine Research Institute)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

use strict;
use warnings;
use Getopt::Long;

# pipe script for analysing bedtools genomeCov output

my $incov;
my $lowc;
my $midc;
my $highc;
my $outcsv = "coverage_stats.csv";
my $junk = 80;
my $suspect = 80;

my $usage = "
USAGE:
purge_haplotigs  cov  -i aligned.bam.genecov  -l <integer>  -m <integer>  -h <integer>  [-o coverage_stats.csv -j $junk  -s $suspect ]

REQUIRED:
-i / -in        The bedtools genomecov output that was produced from 'purge_haplotigs readhist'
-l / -low       The read depth low cutoff (use the histogram to eyeball these cutoffs)
-h / -high      The read depth high cutoff
-m / -mid       The low point between the haploid and diploid peaks

OPTIONAL:
-o / -out       Choose an output file name (CSV format, DEFAULT = $outcsv)
-j / -junk      Auto-assign contig as \"j\" (junk) if this percentage or greater of the contig is 
                low/high coverage (DEFAULT = $junk, > 100 = don't junk anything)
-s / -suspect   Auto-assign contig as \"s\" (suspected haplotig) if this percentage or less of the
                contig is diploid level of coverage (DEFAULT = $suspect)

";


GetOptions(
    "in=s" => \$incov,
    "out=s" => \$outcsv,
    "low=i" => \$lowc,
    "high=i" => \$highc,
    "mid=i" => \$midc,
    "suspect=i" => \$suspect,
    "junk=i" => \$junk
) or die $usage;

($incov) && ($outcsv) && ($lowc || $lowc == 0) && ($highc) && ($midc) or die $usage;

my %out;    # $out{contig}{n_low/n_hap/etc}

# read genome coverage
open my $IN, '<', $incov or die "failed to open $incov for reading\n";
while(<$IN>){
    my@l=split/\s+/;
    
    # ignore genome histogram
    next if ($l[0] eq 'genome');
    
    my $c = \%{$out{$l[0]}};
    
    # add to counts
    if ($l[1] < $lowc) {
        $$c{low} += $l[2]
    } elsif ($l[1] >= $lowc && $l[1] <= $midc) {
        $$c{hap} += $l[2]
    } elsif ($l[1] > $midc && $l[1] <= $highc) {
        $$c{dip} += $l[2]
    } else {
        $$c{high} += $l[2]
    }
}
close$IN;

# rank and print contig stats
open my $OUT, ">", $outcsv or die "failed to open $outcsv for writing\n";

# print header
print $OUT "#contig,contig_reassign,bases_hap_dip,bases_low_high,bases_all,perc_low_coverage,perc_hap_coverage,perc_dip_coverage,perc_high_coverage\n";

# print name-sorted contigs
for my $contig (sort keys %out){
    print_contig($contig);
}
close$OUT;

# done
print STDERR "
Analysis finished successfully! Contig coverage stats saved to '$outcsv'. 
";

exit(0);


sub print_contig {
    my $contig = $_[0];
    my $c = \%{$out{$contig}};
    
    # fill empty values
    for my $i (qw/low hap dip high/){
        $$c{$i} or $$c{$i}=0;
    }
    
    # calc percentages
    my $base_count = $$c{low} + $$c{hap} + $$c{dip} + $$c{high};
    
    # catch instances where no mapping at all to contig and flag as junk (if junking)
    if ($base_count==0){
        $base_count = 1;
        $$c{low} = 1;
    }
    
    my $low_p = sprintf("%.3f", (100 * ($$c{low} / $base_count)));
    my $hap_p = sprintf("%.3f", (100 * ($$c{hap} / $base_count)));
    my $dip_p = sprintf("%.3f", (100 * ($$c{dip} / $base_count)));
    my $high_p = sprintf("%.3f", (100 * ($$c{high} / $base_count)));
    
    my $b_hap_dip = $$c{hap} + $$c{dip};
    my $b_low_high = $$c{low} + $$c{high};
    
    # assign contig
    my $assign = "";
    if ( ($low_p + $high_p) >= $junk ){
        $assign = "j";
    } elsif ($dip_p <= $suspect){
        $assign = "s";
    }
    
    # print
    print $OUT "$contig,$assign,$b_hap_dip,$b_low_high,$base_count,$low_p,$hap_p,$dip_p,$high_p\n";
    
    return;
}


