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


# Adaptation of nucmer2ncbiPlacement.py from https://github.com/skingan/NCBI_DiploidAssembly/blob/master/nucmer2ncbiPlacement.py
# The adaptation is necessary if the assembly has been curated with purge_haplotigs as the FALCON IDs won't necessarily match up.


# PIPELINE:
    # blastn haplotigs against primary contigs and get best hit for each
    # nucmer, delta-filter, show-coords each haplotig against it's blastn hit
    # get the alignment information from the coords file and write to the output placements file



#---GLOBAL VARIABLES---
my $primary;
my $haplotigs;
my $out = "ncbi_placements.tsv";
my $threads = 4;
my $coverage = 50;


my %output :shared;     # $output{$htig}{len} = h-tig length
                        #               {top} = top hit
                        #               {rlen} = ref alignments running total (negative = rev alignments)
                        #               {hlen} = htig alignments running total
                        #               {hs} = htig start
                        #               {he} = htig end
                        #               {rs} = ref start
                        #               {re} = ref end

my $usage = "
USAGE:
purge_haplotigs  ncbiplace  -p primary_contigs.fasta  -h haplotigs.fasta  [ -o outfile  -t threads  -c coverage]

REQUIRED:
-p / -primary       Primary contigs fasta file
-h / -haplotigs     Haplotigs fasta file

OPTIONAL:
-o / -out           Out file name (DEFAULT = $out)
-t / -threads       Threads for blastn and nucmer steps (DEFAULT = $threads)
-c / -coverage      Coverage cutoff percentage for pairing contigs (DEFAULT = $coverage %)

IMPORTANT NOTES:  
    Will only work with nucmer v3.9 or higher

    The placement file generated is currently untested and this script is still under development.
    This script is intended to be an adaptation of nucmer2ncbiPlacement.py from:
    https://github.com/skingan/NCBI_DiploidAssembly/blob/master/nucmer2ncbiPlacement.py
    The adaptation is necessary if the assembly has been curated with Purge Haplotigs as the 
    curated primary contigs and haplotigs wont necessarily be paired by the FALCON Unzip IDs.

";


#---PREFLIGHT CHECKS---
check_programs("minimap2", "samtools", "nucmer", "delta-filter", "show-coords") or err("One or more dependencies missing, exiting");



#---PARSE ARGS---
my $args = "@ARGV";
GetOptions (
    "primary=s" => \$primary,
    "haplotigs=s" => \$haplotigs,
    "out=s" => \$out,
    "threads=i" => \$threads,
    "coverage=i" => \$coverage
) or die $usage;


# check args and files
($primary) && ($haplotigs) || die $usage;

check_files($primary, $haplotigs);




#---SET UP LOGGING---
our $LOG;
my $TMP = "tmp_purge_haplotigs/NCBIPLACE";
(-d $_) or mkdir $_ for ("tmp_purge_haplotigs", $TMP);

open $LOG, ">", "tmp_purge_haplotigs/purge_haplotigs_ncbiplace.log" or err("failed to open log file for writing");



                # clean file names etc
                my $primary_fname = $primary;
                my $haplotigs_fname = $haplotigs;
                ($_ =~ s/.*\///) for ($primary_fname, $haplotigs_fname);



# other misc setup
my $minimap2_parameters = "-t $threads";
my $MINIMAP2_OUT = "$TMP/minimap2.paf.gz";



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



# print settings
msg("
Beginning Pipeline

PARAMETERS:
Primary contigs FASTA       $primary
Haplotigs FASTA             $haplotigs
Out FASTA                   $out
Threads                     $threads
Coverage align cutoff       $coverage %

RUNNING USING COMMAND:
purge_haplotigs ncbiplace $args
");


# get fasta indexes ready and get the haplotig IDs and lengths
((-s "$_.fai") || runcmd({ command => "samtools faidx $_ 2> $TMP/samtools_faidx.stderr", 
                           logfile => "$TMP/samtools_faidx.stderr" })) for ($haplotigs, $primary);



# read in the contig lengths
read_contig_lens();

# run minimap2
run_minimap2();

# get hits from minimap2 output
get_minimap2_hits();


# open outfile for writing
open my $OUT, ">", $out or err("failed to open $out for writing");
# print the output header
print $OUT "#alt_asm_name\tprim_asm_name\talt_scaf_name\tparent_type\tparent_name\tori\talt_scaf_start\talt_scaf_stop\tparent_start\tparent_stop\talt_start_tail\talt_stop_tail\n";


#---NUCMER TO CHARACTERIZE ALIGNMENTS---
nucmer_alignments();






#---WRITE THE PLACEMENT FILE---
for my $htig(sort(keys(%output))){
    my $h = $output{$htig};
    (($$h{hlen}) && ($$h{hlen} > (($coverage / 100) * $$h{len}))) ?                     # write the placement if > alignment cutoff to the best hit primary contig
        (print $OUT "Haplotigs\tPrimary_Assembly\t$htig\tSCAFFOLD\t$$h{top}\t",         # contig/assembly names
        ($$h{rlen} < 0 ? "-" : "+"),                                                    # check orientation (ncbi's 'mixed' tag not implemented)
        "\t$$h{hs}\t$$h{he}\t$$h{rs}\t$$h{re}\t",                                       # start/stop positions
        ($$h{hs} == 1 ? 0 : ($$h{hs}-1)), "\t",                                         # check start tail
        ($$h{he} == $$h{len} ? 0 : ($$h{len} - $$h{he})), "\n")                         # check end tail
    : print $OUT "Haplotigs\tPrimary_Assembly\t$htig", ("\tna" x 9), "\n";          # no placements if no or poor alignment
}

# DONE.
exit(0);



sub read_contig_lens {
    open my $HFI, "<", "$haplotigs.fai" or err("failed to open file $haplotigs.fai");
    while(<$HFI>){
        my @l = split(/\s+/, $_);
        my %l : shared;
        $output{$l[0]} = \%l;
        $output{$l[0]}{len} = $l[1];
    }
    close $HFI;
    return;
}


sub run_minimap2 {
    if (-s $MINIMAP2_OUT){
        msg("Minimap2 alignments file found, skipping minimap2 hit search");
    } else {
        msg("Running minimap2 hit search");
        runcmd({ command => "minimap2 $minimap2_parameters $primary $haplotigs 2> $TMP/minimap2.stderr | gzip - > $MINIMAP2_OUT.tmp",
                 logfile => "$TMP/minimap2.stderr",
                 silent => 1 });
        rename "$MINIMAP2_OUT.tmp", "$MINIMAP2_OUT";
        msg("Minimap2 hit search done");
    }
    return;
}


sub get_minimap2_hits {
    msg("reading and sorting minimap2 hits");
    my %hits;
    open my $PAF, "zcat $MINIMAP2_OUT |" or err("failed to open pipe \"zcat $MINIMAP2_OUT |\"");
    while(<$PAF>){
        my@l=split/\s+/;
        if ($l[0] ne $l[5]){
            $hits{$l[0]}{$l[5]}+=$l[9];
        }
    }
    close$PAF;
    for my $haplotig (keys %hits){
        for my $hit (sort { $hits{$haplotig}{$b} <=> $hits{$haplotig}{$a} } keys %{$hits{$haplotig}} ){
            $output{$haplotig}{top} = $hit;
            last;
        }
    }
    return;
}


sub nucmer_alignments {
    msg("characterizing alignments with nucmer");
    
    # enqueue nucmer jobs
    my $job=0;
    for my $htig(sort(keys(%output))){
        my $htig_out = $htig;
        $htig_out =~ s/\|.+//;       # remove |quiver|arrow|etc tags
        my $h = $output{$htig};
        if ($$h{top}){
            my $pctg = $$h{top};
            $pctg =~ s/\|.+//;
            $nucQueue->enqueue([$htig, $pctg, $htig_out, $job]);
            $job++;
        }
    }
    # finalise
    $nucQueue->end();
    
    # spawn workers
    for (1..$threads){
        $available_threads->down(1);
        threads->create(\&nuc_align);
    }
    
    # wait for workers to finish
    $available_threads->down($threads);
    $available_threads->up($threads);
    
    # join workers
    for my $thr (threads->list()){
        $thr->join();
    }
    
    return;
}


sub nuc_align {
    
    while (defined(my $args = $nucQueue->dequeue())) {
    
        my $htig = @{$args}[0];
        my $pctg = @{$args}[1];
        my $htig_out = @{$args}[2];
        my $job = @{$args}[3];
        
        my $h = $output{$htig};
        
        msg("processing haplotig $htig -> $$h{top}");
        # run nucmer etc if needed
        if (!(-s "$TMP/$htig_out.$pctg.coords")){
            qruncmd("samtools faidx $haplotigs \"$htig\" > $TMP/$job.htig.fa");
            qruncmd("samtools faidx $primary \"$$h{top}\" > $TMP/$job.ref.fa");
            qruncmd("nucmer $TMP/$job.htig.fa $TMP/$job.ref.fa -p $TMP/$job.tmp");
            qruncmd("delta-filter -g $TMP/$job.tmp.delta > $TMP/$job.tmp.g.delta");
            qruncmd("show-coords -rT $TMP/$job.tmp.g.delta > $TMP/$htig_out.$pctg.coords.tmp");
            qruncmd("mv $TMP/$htig_out.$pctg.coords.tmp $TMP/$htig_out.$pctg.coords");
        }

        # parse the coords file
        open my $COR, "$TMP/$htig_out.$pctg.coords" or err("failed to open $TMP/$htig_out.$pctg.coords");
        $writing_to_out->down(1);
        $$h{hs} = 999999999;
        $$h{rs} = 999999999;
        $$h{he} = 0;
        $$h{re} = 0;
        while(<$COR>){
            next if ($_ !~ /^\d/);
            my @l = split(/\s+/, $_);
            $$h{rlen} += $l[3] - $l[2];
            $$h{hlen} += $l[1] - $l[0];
            $$h{rs} < min($l[2],$l[3]) or $$h{rs} = min($l[2],$l[3]);
            $$h{re} > max($l[2],$l[3]) or $$h{re} = max($l[2],$l[3]);
            $$h{hs} < $l[0] or $$h{hs} = $l[0];
            $$h{he} > $l[1] or $$h{he} = $l[1];
        }
        $writing_to_out->up(1);
        close $COR;
        
        # cleanup and reset
        (unlink $_) for ("$TMP/$job.htig.fa", "$TMP/$job.ref.fa", "$TMP/$job.tmp.g.delta", "$TMP/$job.tmp.delta");
    }
    
    # close worker
    $available_threads->up(1);
    
    return;
}



