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


# Input parameters
my $primaryFasta;
my $haplotigsFasta;
my $minLen = 10000;
my $gapLen = 10000;
my $gapRatio = 0.5;
my $threads = 4;
my $mm2Drop = '4G';
my $outPrefix = 'clip';


my $usage = "
USAGE:
purge_haplotigs  clip  -p contigs.fasta  -h haplotigs.fasta

REQUIRED:
-p / -primary       Primary contigs (curated.fasta by default) from the 'purge' stage.
-h / -haplotigs     Haplotigs (curated.haplotigs.fasta by default) from the 'purge' stage.

OPTIONAL:
-o / -outPrefix     Output file prefix. DEFAULT = $outPrefix
-l / -length        Minimum overlap length to trigger clipping. DEFAULT = $minLen
-g / -gap           Maximum end gap to allow clipping. DEFAULT = $gapLen
-r / -ratio         Max allowable gap to alignment ratio. DEFAULT = $gapRatio
-t / -threads       Threds to use. DEFAULT = $threads

NOTE: This script is experimental. 
";


# Parse arguments
my $args = "@ARGV";
GetOptions (
    "primary=s" => \$primaryFasta,
    "haplotigs=s" => \$haplotigsFasta,
    "outPrefix=s" => \$outPrefix,
    "length=i" => \$minLen,
    "ratio=f" => \$gapRatio,
    "threads=i" => \$threads,
    "gap=i" => \$gapLen
) or die $usage;

($primaryFasta) && ($haplotigsFasta) || die $usage;


# minimap run parameters
my $mm2Parm = "-f 0.001 -X -m 500 -g $minLen -r $minLen";


# Set up logging
our $LOG;
my $TMP_DIR = 'tmp_purge_haplotigs';
(-d $TMP_DIR) or mkdir $TMP_DIR;
open $LOG, '>', "$TMP_DIR/purge_haplotigs_clip.log" or die 'failed to open log file for writing';


# Directories
my $TMP = "$TMP_DIR/CLIP";
(-d $TMP) or mkdir $TMP;
my $ERR = "$TMP/STDERR";
(-d $ERR) or mkdir $ERR;


# Files
my $prefix = $primaryFasta;
$prefix =~ s/.+\///;
$prefix =~ s/\.f[nast]+$//i;
my $mm2Paf = "$TMP/$prefix.$mm2Drop.paf";
my $mm2Idx = "$mm2Paf.index";
my $hBed = "$TMP/clip.bed";
my $pBed = "$TMP/keep.bed";


# Global vars
my %mmIndex;                # $mmIndex{contig} = byte file read start position for contig
my %ctgLen;                 # $ctgLen{contig} = length of contig
my @clip;                   # nested array ( (ctg, hStart, hStop),...)
my %clp;                    # quick lookup index for clipped contigs

# Pre flight
if (!(check_programs qw/minimap2 samtools bedtools/)){
    err('ONE OR MORE REQUIRED PROGRAMS IS MISSING');
}


# print params
msg("\n
PARAMETERS:
Primary Contigs         $primaryFasta
Haplotigs               $haplotigsFasta
Min length to clip      $minLen
Threads                 $threads
Minimap2:
- minimiser drop        $mm2Drop
- alignment params      $mm2Parm

RUN COMMAND:
purge_haplotigs clip $args
");



# SCRIPT START

# slurp contig lengths
slurpLens();

# All v All align the primary contigs
primaryContigAlign();

# Read the index of alignments
readAlignmentsIndex();

# find overlaps and trim
findEdgeOverlaps();

# check of overlaps of clipping coords and finalize
mergeClipping();

# print output
writeOutput();

# done
exit(0);



sub slurpLens {
    msg('Reading genome indexes');
    # index fasta if needed
    if (!(-s "$primaryFasta.fai")){
        runcmd({
            command => "samtools faidx $primaryFasta 2> $ERR/samtools.stderr",
            logfile => "$ERR/samtools.stderr",
            silent => 1
        })
    }
    
    # read lengths
    open my $FAI, '<', "$primaryFasta.fai" or err("failed to open $primaryFasta.fai for reading");
    while(<$FAI>){
        my@l=split/\s+/;
        $ctgLen{$l[0]}=$l[1];
    }
    close $FAI;
    
    return;
}


sub primaryContigAlign {
    
    # clean up old files
    (-e "$mm2Paf.tmp") and unlink "$mm2Paf.tmp";
    (-e "$mm2Idx.tmp") and unlink "$mm2Idx.tmp";
    
    # run the alignments if needed
    if (!(-s $mm2Paf)){
        msg('Running minimap2');
        runcmd ({ 
            command => "minimap2 $mm2Parm -I $mm2Drop -t $threads $primaryFasta $primaryFasta 2> $ERR/minimap2.stderr | sort -k1,1 -k6,6 -k3,3n > $mm2Paf.tmp",
            logfile => "$ERR/minimap2.stderr",
            silent => 1
        });
        rename "$mm2Paf.tmp", $mm2Paf;
    } else {
        msg('Skip minimap2 alignments');
    }
    
    
    # index the alignments if needed
    if (!(-s $mm2Idx)){
        
        msg('Indexing minimap2 alignments');
        
        # open the alignments for reading
        open my $FWR, '<', $mm2Paf or err("failed to open $mm2Paf for reading");
        
        # open filehandle for alignment index
        open my $IDX, '>', "$mm2Idx.tmp" or err("failed to open $mm2Idx.tmp for writing");
        
        # tracking for index file
        my $curCtg = 'init';
        my $offset = tell($FWR);
        
        # generate the index file
        while(<$FWR>){
            my ($ctg) = $_ =~ m/^(.+?)\s/;
            if ($curCtg ne $ctg){
                print $IDX $ctg, "\t", $offset, "\n";
                $curCtg = $ctg;
            }
            $offset = tell($FWR);
        }
        close $FWR;
        close $IDX;
        
        # finalize the index file
        rename "$mm2Idx.tmp", $mm2Idx;
    } else {
        msg('Skip indexing minimap2 alignments');
    }
    
    return;
}


sub readAlignmentsIndex {
    msg('Reading index of minimap2 alignments');
    open my $IDX, '<', $mm2Idx or err("failed to open $mm2Idx for reading");
    while(<$IDX>){
        my@l=split/\s+/;
        $mmIndex{$l[0]} = $l[1];
    }
    close $IDX;
    
    return;
}


sub findEdgeOverlaps {
    msg('Finding edge overlaps');
    # iterate each contig, smallest to largest
    for my $ctg ( sort { $ctgLen{$a} <=> $ctgLen{$b} } keys %mmIndex){
        
        # slurp all the alignments from the PAF for contig
        my @ali;     # nested array ( (ref, rLen, start, stop, query, qLen, start, stop), ...)
        @ali = getAlignments($ctg);
        
        # iterate alignments
        for my $i (0..$#ali){
            $a = \@{$ali[$i]};
            
            # shuffle so that the shorter contig is 0-3 and the longer 4-7
            if ($$a[1] > $$a[5]){
                @{$ali[$i]} = ($$a[4], $$a[5], $$a[6], $$a[7], $$a[0], $$a[1], $$a[2], $$a[3]);
            }
            
            # skip alignments with length less than the cutoff
            next if ($$a[3] - $$a[2] < $minLen);
            
            # find alignments on end of reference (shorter contig for clipping)
            if ( ($$a[2] < $gapLen) or ($$a[3] > $$a[1] - $gapLen) ){
                
                # find on end of query (longer contig, no clipping)
                if ( ($$a[6] < $gapLen) or ($$a[7] > $$a[5] - $gapLen) ){
                    
                    # check gap to alignment ratio
                    next if ( ($$a[2] < $gapLen) and ($$a[2] > $gapRatio * ($$a[3] - $$a[2])) );
                    next if ( ($$a[3] > $$a[1] - $gapLen) and ($$a[1] - $$a[3] > $gapRatio * ($$a[3] - $$a[2])) );

                    # extend the coords to the contig ends for clipping
                    ($$a[2] < $gapLen) and $$a[2] = 0;
                    ($$a[3] > $$a[1] - $gapLen) and $$a[3] = $$a[1];

                    # push the coords for clipping
                    @{$clip[@clip]} = ($$a[0], $$a[2], $$a[3]);
                    
                    # quick lookup index of contigs flagged for clipping
                    $clp{$$a[0]}=1;
                    
                    msg("[CLIP] $$a[0]:$$a[2]-$$a[3] --> [KEEP] $$a[4]:$$a[6]-$$a[7]");
                }
            }
        }
    }
    
    # Check if no clipping needed
    if (@clip == 0){
        msg('No overlapping edges detected, no clipping needed');
        exit(0);
    }
    
    return;
}


sub getAlignments {
    my $ctg = $_[0];
    my @ali;
    
    # open the PAF
    open my $FWR, '<', $mm2Paf or err("failed to open $mm2Paf for reading");
    
    # fast-forward to to the correct location in the alignment file
    seek($FWR, $mmIndex{$ctg}, 0);
    
    # grab the alignments
    while(<$FWR>){
        my@l=split/\s+/;
        last if ($l[0] ne $ctg);
        if ($l[0] ne $l[5]){
            #                ref    rLen   start  stop   query  qLen   start  stop
            @{$ali[@ali]} = ($l[0], $l[1], $l[2], $l[3], $l[5], $l[6], $l[7], $l[8]);
        }
    }
    
    return @ali;
}


sub mergeClipping {
    msg('Merge overlapping clipping coords');
    
    # dump coords to bed file
    open my $CLP, '>', $hBed or err("failed to open $hBed for writing");
    for my $i (0..$#clip){
        print $CLP "@{$clip[$i]}[0]\t@{$clip[$i]}[1]\t@{$clip[$i]}[2]\n";
    }
    close $CLP;
    
    # merge overlaps
    runcmd({
        command => "cat $hBed | sort -k1,1 -k2,2n | bedtools merge 1> $hBed.tmp 2> $ERR/bedtools.stderr",
        logfile => "$ERR/bedtools.stderr",
        silent => 1
    });
    
    # overwrite bed file
    rename "$hBed.tmp", $hBed;
    
    # gen file of contigs for clipping for bedtools
    my $gen = "$TMP/clip.gen";
    open my $GEN, '>', $gen or err("failed to open $gen for writing");
    for my $ctg (sort keys %clp){
        print $GEN "$ctg\t$ctgLen{$ctg}\n";
    }
    close $GEN;
    
    # bedtools to get the regions to keep for each contig being clipped
    runcmd({
        command => "bedtools complement -i $hBed -g $gen 1> $pBed 2> $ERR/bedtools.stderr",
        logfile => "$ERR/bedtools.stderr",
        silent => 1
    });
    
    return;
}


sub writeOutput {
    msg('Writing new assemblies');
    
    my $primary = "$outPrefix.fasta";
    my $haplotigs = "$outPrefix.haplotigs.fasta";
    
    (-s $primary) and unlink $primary;
    (-s $haplotigs) and unlink $haplotigs;
    
    # slurp the clipping regions
    my %keep;
    my @trim;
    
    open my $KEP, '<', $pBed or err("failed to open $pBed for reading");
    while(<$KEP>){
        my@l=split/\s+/;
        $keep{$l[0]}{A} = $l[1];
        $keep{$l[0]}{B} = $l[2];
    }
    close $KEP;
    
    open my $CLP, '<', $hBed or err("failed to open $hBed for reading");
    while(<$CLP>){
        my@l=split/\s+/;
        @{$trim[@trim]} = ($l[0],$l[1],$l[2]);
    }
    close $CLP;
    
    # write the primary assembly
    for my $ctg (sort { $ctgLen{$b} <=> $ctgLen{$a} } keys %ctgLen){
        my $cmd;
        if ($clp{$ctg}){
            if (!($keep{$ctg})){
                next;
            }
            $cmd = "samtools faidx $primaryFasta \'$ctg\':$keep{$ctg}{A}-$keep{$ctg}{B}";
        } else {
            $cmd = "samtools faidx $primaryFasta \'$ctg\'";
        }
        runcmd({
            command => "$cmd 1>> $primary 2> $ERR/samtools.stderr",
            logfile => "$ERR/samtools.stderr",
            silent => 1 });
    }
    
    # write the new haplotigs
    runcmd ({
        command => "cat $haplotigsFasta > $haplotigs",
        silent => 1 });
    
    for my $i (0..$#trim){
        my $t = \@{$trim[$i]};
        my $cmd = "samtools faidx $primaryFasta \'$$t[0]\':$$t[1]-$$t[2]";
        runcmd ({
            command => "$cmd 1>> $haplotigs 2> $ERR/samtools.stderr",
            logfile => "$ERR/samtools.stderr",
            silent => 1 });
    }
    
    return;
}











