#!/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;
use FindBin qw($RealBin);
use threads;
use Thread::Semaphore;
use Thread::Queue;
use lib "$RealBin/../lib";
use PipeUtils;



# pre flight
if (check_programs("bedtools", "Rscript", "samtools")){
    msg("ALL DEPENDENCIES OK\n\n");
} else {
    err("ONE OR MORE DEPENDENCIES MISSING");
}
my $SCRIPT = "$RealBin/../scripts/";



# vars
my $genome;
my $bamfile;
my $threads = 4;
my %contigs;
my %genome_coverage;


# help msg
my $usage = "
USAGE:
purge_haplotigs  readhist  -b aligned.bam  -g genome.fasta  [ -t threads ]

REQUIRED:
-b / -bam       BAM file of aligned and sorted reads/subreads to the reference
-g / -genome    Reference FASTA for the BAM file.

OPTIONAL:
-t / -threads   Number of worker threads to use, DEFAULT = $threads, MINIMUM = 2

";



# parse args
GetOptions (
    "genome=s" => \$genome,
    "bam=s" => \$bamfile,
    "threads=i" => \$threads
) or die $usage;



# check args and files
($genome) && ($bamfile) or die $usage;
my $bamfilename = $bamfile;
$bamfilename =~ s/.*\///;   # remove file path

if (!(check_files($bamfile, $genome))){
    die $usage;
}



#---THREADS---
my $available_threads = Thread::Semaphore->new($threads);
my $writing_to_out = Thread::Semaphore->new(1);
my $queue = Thread::Queue->new();



#---LOGGING---
our $LOG;
my $TMP_DIR = "tmp_purge_haplotigs";
my $ERR_DIR = "$TMP_DIR/STDERR";
my $MISC_DIR = "$TMP_DIR/MISC";
((-d $_) || mkdir $_) for ($TMP_DIR, $ERR_DIR, $MISC_DIR);
open $LOG, ">", "$TMP_DIR/purge_haplotigs_readhist.log" or die "failed to open log file for writing";



# FILENAMES FOR CLARITY
my $gencov_file = "$bamfilename.gencov";
my $csv_file = "$MISC_DIR/$bamfilename.histogram.csv";
my $png_file = "$bamfilename.histogram.png";



# begin
msg("Beginning read-depth histogram generation");

# read contigs from fasta fai
read_fasta_fai();

# run bedtools genomecov if needed
bedtools_multicov();

# make the histogram csv and png if needed
histo_csv_png();

# done
msg("

READ HISTOGRAM GENERATION FINISHED SUCCESSFULLY! 

Check '$bamfile.histogram.png' to observe where your haploid and diploid peaks are
and choose your low, midpoint, and high cutoffs (check the example histogram png 
in the readme). You will need '$bamfile.gencov' and the cutoffs for the next step 
'purge_haplotigs contigcov'
");

exit(0);



sub read_fasta_fai {
    if (!(-s "$genome.fai")){
        runcmd({ command => "samtools faidx $genome 2> $ERR_DIR/samtools.faidx.stderr",
                 logfile => "$ERR_DIR/samtools.faidx.stderr",
                 silent => 1 });
    }
    open my $FAI, '<', "$genome.fai" or err("failed to open $genome.fai for reading");
    while(<$FAI>){
        my@l=split/\s+/;
        $contigs{$l[0]} = $l[1];
    }
    close $FAI;
    
    return;
}


sub bedtools_multicov {
    if (!(-s $gencov_file)){
        msg("running genome coverage analysis on $bamfile");
        
        # cleanup
        if (-s "$gencov_file.tmp"){
            unlink "$gencov_file.tmp";
        }
        
        # check bam is indexed
        if (!(-s "$bamfile.bai")){
            runcmd ({ command => "samtools index $bamfile 2> $ERR_DIR/samtools.index.stderr",
                      logfile => "$ERR_DIR/samtools.index.stderr",
                      silent => 1 });
        }
        
        # queue the jobs
        for my $contig (sort { $contigs{$b} <=> $contigs{$a} } keys %contigs){
            $queue->enqueue($contig);
        }
        $queue->end();
        
        # spawn worker threads
        for (1..$threads){
            $available_threads->down(1);
            threads->create(\&gencov_job);
        }
        
        # wait on threads and then join
        $available_threads->down($threads);
        $available_threads->up($threads);
        ($_->join()) for threads->list();
        
        # build the genome histogram
        build_gen_histo();
        
    } else {
        msg("$gencov_file found, skipping bedtools genomecov step");
    }
    
    return;
}



sub gencov_job {
    while (defined(my $contig = $queue->dequeue())){
        
        my %out;
        
        # open a pipe and build the histogram
        open my $TMP, "samtools mpileup -r \"$contig\" -f $genome $bamfile 2>> $ERR_DIR/mpileup.stderr | awk '{ print \$4 }' |" or err("failed to open pipe \"samtools mpileup -r \"$contig\" -d 200 -f $genome $bamfile | awk '{ print \$4 }' |\"");
        while(<$TMP>){
            $_=~s/\s//g;
            ($_<=200) or ($_=200);
            $out{$_}++;
        }
        close $TMP;
        
        # print the output
        $writing_to_out->down(1);
        open my $TO, '>>', "$gencov_file.tmp" or err("failed to open $gencov_file.tmp for appending");
        for my $depth ( sort { $a <=> $b } keys %out){
            print $TO "$contig\t$depth\t$out{$depth}\n";
        }
        close $TO;
        $writing_to_out->up(1);
    }
    
    # exit thread
    $available_threads->up(1);

    return;
}



sub build_gen_histo {
    open my $GCF, '<', "$gencov_file.tmp" or err("failed to open $gencov_file.tmp for reading");
    while(<$GCF>){
        my@l=split/\s+/;
        $genome_coverage{$l[1]}+=$l[2];
    }
    close $GCF;

    # append the genome hist to the gencov file
    open $GCF, '>>', "$gencov_file.tmp" or err("failed to open $gencov_file.tmp for appending");
    for my $depth (sort { $a <=> $b } keys %genome_coverage){
        print $GCF "genome\t$depth\t$genome_coverage{$depth}\n";
    }
    close $GCF;
    qruncmd("mv $gencov_file.tmp $gencov_file");
    
    return;
}



sub histo_csv_png {
    # csv
    if (!(-s $csv_file)){
        msg("generating histogram");
        runcmd( { command => "grep genome $gencov_file | awk '{ print \$2 \",\" \$3 }' > $csv_file",
                  silent => 1 } );
    } else {
        msg("$csv_file found, skipping step");
    }

    # png
    if (!(-s $png_file)){
        msg("generating png image");
        runcmd( { command => "$SCRIPT/gen_histogram.Rscript $csv_file $png_file 2> $ERR_DIR/gen_histogram.stderr",
                  logfile => "$ERR_DIR/gen_histogram.stderr",
                  silent => 1 });
    } else {
        msg("$png_file found, skipping Rscript step");
    }
    
    return;
}


