#!/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/";
my $xvfb;
if (system("Rscript -e 'png()' 2> /dev/null")!=0){
    if (check_programs("xvfb-run")){
        $xvfb = "xvfb-run";
    } else {
        msg('WARNING: Neither X11 nor xvfb is available');
    }
} else {
    $xvfb=" ";
}


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



# help msg
my $usage = "
USAGE:
purge_haplotigs  hist  -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
-d / -depth     Maximum cutoff for depth. DEFAULT = $depth, increase if needed,
                set much higher than your expected average coverage.
";



# parse args
GetOptions (
    "genome=s" => \$genome,
    "bam=s" => \$bamfile,
    "threads=i" => \$threads,
    "depth=i" => \$depth
) 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_hist.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
coverage_analysis();

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

# check
if (!(-s $png_file)){
    msg("

WARNING: 
PNG image generation failed. You will need to manually generate a histogram
to continue with the next step using the file: $csv_file
");
} else {
    msg("

Pipeline finished! Your histogram is saved to: $png_file
");
}

# done
msg("
Check your histogram 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 cov'
");


# done
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 coverage_analysis {
    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
        msg('Building 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;
        ($out{$_}=0)for(0..$depth);
        
        # open a pipe and build the histogram
        my $errLog = "$ERR_DIR/samtoolsDepth.stderr";
        my $cmd = "samtools depth -a -d $depth -r \"$contig\" $bamfile 2>> $errLog";
        open my $TMP, '-|', $cmd or err("failed to open pipe: $cmd | (this script), check $errLog");
        
        while(<$TMP>){
            my@l=split/\s+/;
            $out{$l[2]}++;
        }
        close $TMP or err("Failed to close pipe: $cmd | (this script), check $errLog");
        
        # 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;
    rename "$gencov_file.tmp", $gencov_file;
    
    return;
}



sub histo_csv_png {
    # csv
    my %hist;
    msg("generating histogram");
    open my $TMP, '<', $gencov_file or die;
    while(<$TMP>){
        my@l=split/\s+/;
        if($l[0] eq 'genome'){
            if ($l[1]>$depth){
                $hist{$depth} += $l[2];
            } else {
                $hist{$l[1]} = $l[2];
            }
        }
    }
    close$TMP;
    
    open my $CSV, '>', $csv_file or die;
    for my $d (sort { $a <=> $b } keys %hist){
        print $CSV $d, ",", $hist{$d}, "\n";
    }
    close$CSV;

    # png
    if ($xvfb){
        msg("generating png image");
        runcmd( { command => "$xvfb $SCRIPT/gen_histogram.Rscript $csv_file $png_file $depth 2> $ERR_DIR/gen_histogram.stderr",
                  logfile => "$ERR_DIR/gen_histogram.stderr",
                  silent => 1 });
    }
    
    return;
}


