#!/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 lib "$RealBin/../lib";
use PipeUtils;


# 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.


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

my %output;             # $output{$htig}{top} = top hit
                        #               {sf} = sum forward alignments
                        #               {sr} = sum reverse alignments
                        #               {hs} = htig start
                        #               {he} = htig end
                        #               {rs} = ref start
                        #               {re} = ref end
my %prim;
my %hap;
my %falcCount;
$falcCount{'9'x6 . 'F'} = 0;        # prefix for unplaced haploitgs
my %falcName;



my $usage = "
USAGE:
purge_haplotigs  place  -p primary_contigs.fasta  -h haplotigs.fasta  \
            [ -o out.tsv  -t <INT  -c <INT  -f]

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

OPTIONAL:
-o / -out           Placement file name. DEFAULT = $out
-t / -threads       Threads for Minimap2. DEFAULT = $threads
-c / -coverage      Coverage cutoff percentage for pairing contigs. 
                    DEFAULT = $coverage (%)
-f / -falconNaming  Rename contigs in the style used by FALCON Unzip. Saved
                    to <in-prefix>.FALC.fasta

";


#---PREFLIGHT CHECKS---
check_programs(qw/minimap2 samtools/) 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,
    "falconNaming" => \$falconNaming
) 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/PLACE";
(-d $_) or mkdir $_ for ("tmp_purge_haplotigs", $TMP);

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



# clean file names etc
my $pName = $primary;
my $hName = $haplotigs;
($_ =~ s/.*\///) for ($pName, $hName);
($_ =~ s/\.f[nast]+$//) for ($pName, $hName);



# other misc setup
my $minimap2_parameters = "-t $threads --secondary=no -r 10000";
my $MINIMAP2_OUT = "$TMP/$pName-$hName.paf.gz";



# 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 place $args
");


# contig lengths
contig_lengths();

# run minimap2
run_minimap2();

# get hits from minimap2 output
get_minimap2_hits();

# get the minimap2 alignment coords for the top hits
get_alignment_coords();

# rename the contigs if flagged
($falconNaming) and falcon_name();

# write the output
write_placement();


# DONE.
msg('Done!');
exit(0);


sub contig_lengths {
    msg('Reading contig lengths');
    # get fasta indexes
    for my $fa ($haplotigs, $primary){
        if(!(-s "$fa.fai")){
            msg("Indexing $fa");
            runcmd({ command => "samtools faidx $fa 2> $TMP/samtools_faidx.stderr", 
                     logfile => "$TMP/samtools_faidx.stderr",
                     silent => 1})
        }
    }
    
    # slurp contig names and lengths
    open my $FAI, '<', "$primary.fai" or err("failed to open $primary.fai for reading");
    while(<$FAI>){
        my@l=split/\s+/;
        $prim{$l[0]}=$l[1];
    }
    close $FAI;
    open $FAI, '<', "$haplotigs.fai" or err("failed to open $haplotigs.fai for reading");
    while(<$FAI>){
        my@l=split/\s+/;
        $hap{$l[0]}=$l[1];
    }
    close $FAI;
}


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 minimap2 alignments');
    my %hits;
    open my $PAF, "gunzip -c $MINIMAP2_OUT |" or err("failed to open pipe \"gunzip -c $MINIMAP2_OUT |\"");
    while(<$PAF>){
        my@l=split/\s+/;
        if ($l[0] ne $l[5]){
            $hits{$l[0]}{$l[5]}+=$l[9];
        }
    }
    close$PAF;
    msg('Getting best hits');
    for my $haplotig (keys %hits){
        for my $hit (sort { $hits{$haplotig}{$b} <=> $hits{$haplotig}{$a} } keys %{$hits{$haplotig}} ){
            my $h = \%{$output{$haplotig}};
            # initialize coords
            $$h{top} = $hit;
            $$h{hs} = 999999999;
            $$h{rs} = 999999999;
            $$h{he} = 0;
            $$h{re} = 0;
            $$h{sf} = 0;
            $$h{sr} = 0;
            last;
        }
    }
    return;
}


sub get_alignment_coords {
    msg('Getting alignments coords for best hits');
    
    # read through alignment file and build coords on the fly
    open my $PAF, "gunzip -c $MINIMAP2_OUT |" or err("failed to open pipe \"gunzip -c $MINIMAP2_OUT |\"");
    while(<$PAF>){
        my@l=split/\s+/;
        my $h = \%{$output{$l[0]}};
        if($l[5] eq $$h{top}){
            # direction
            ($l[4]eq'+') ? ($$h{sf}+=$l[10]) : ($$h{sr}+=$l[10]);
            # coords
            $$h{hs} < $l[2] or $$h{hs} = $l[2];
            $$h{he} > $l[3] or $$h{he} = $l[3];
            $$h{rs} < $l[7] or $$h{rs} = $l[7];
            $$h{re} > $l[8] or $$h{re} = $l[8];
        }
    }
    close $PAF;
    return;
}


sub falcon_name {
    msg('Renaming contigs and haplotigs in the FALCON Unzip format');
    
    # generate new primary contig names
    my $p=0;
    for my $pctg (sort { $prim{$b} <=> $prim{$a} } keys %prim){
        $falcName{$pctg} = sprintf("%06d", $p) . "F";
        $p++;
        $falcCount{$falcName{$pctg}} = 0;
    }
    
    # generate new haplotig names
    for my $htig (sort { $hap{$b} <=> $hap{$a} } keys %hap){
        if ($output{$htig}{top}){
            $falcName{$htig} = $falcName{$output{$htig}{top}} . "_" . sprintf("%03d", $falcCount{$falcName{$output{$htig}{top}}});
            $falcCount{$falcName{$output{$htig}{top}}}++;
        } else {
            $falcName{$htig} = '9'x6 . 'F_' . sprintf("%03d", $falcCount{'9'x6 . 'F'});
            $falcCount{'9'x6 . 'F'}++;
        }
    }
    
    # write the new files
    my $outPrimary = "$pName.FALC.fasta";
    my $outHaplotigs = "$hName.FALC.fasta";
    
    renameFasta($primary, $outPrimary);
    renameFasta($haplotigs, $outHaplotigs);
    
    return;
}


sub renameFasta {
    my $in = $_[0];
    my $out = $_[1];
    
    open my $IN, '<', $in or err("failed to open $in for reading");
    open my $OUT, '>', $out or err("failed to open $out for writing");
    
    while(<$IN>){
        if ($_ =~ />(\S+)\s/){
            print $OUT ">$falcName{$1}\n";
        } else {
            print $OUT $_;
        }
    }
    close $IN;
    close $OUT;
    
    return;
}


sub write_placement {
    msg('Writing placement file');
    # 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";
    
    for my $htig(sort(keys(%hap))){
        my $h = \%{$output{$htig}};
        
        # check coverage of haplotig with alignments, shouldn't be an issue.
        if ( !($$h{top}) or ($$h{he} - $$h{hs} < ($coverage / 100) * $hap{$htig}) ){
            print $OUT "Haplotigs\tPrimary_Assembly\t", ($falconNaming ? $falcName{$htig} : $htig),  ("\tna" x 9), "\n";
        } else {
            if ($falconNaming){
                print $OUT "Haplotigs\tPrimary_Assembly\t$falcName{$htig}\tSCAFFOLD\t$falcName{$$h{top}}\t"
            } else {
                print $OUT "Haplotigs\tPrimary_Assembly\t$htig\tSCAFFOLD\t$$h{top}\t"
            }
            
            print $OUT ($$h{sf} > $$h{sr} ? '+' : '-'),                                     # 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} > 0 ? ($$h{hs}-1) : 0), "\t",                                          # check start tail
            ($$h{he} == $hap{$htig} ? 0 : ($hap{$htig} - $$h{he})), "\n";                   # check end tail
        }
    }

    close $OUT;
    
    return;
}






