#!/usr/bin/env python
#MIT License
#
#Copyright (c) 2019 Iñigo Prada Luengo
#
#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.

from __future__ import division


import os
import sys
from Bio.Seq import Seq
import time
from utils import *
import pandas as pd
import traceback



class realignment:
    """Class for managing the realignment and eccDNA indetification of circle-map"""

    def __init__(self, input_bam,qname_bam,sorted_bam,genome_fasta,directory,mapq_cutoff,insert_size_mapq,std_extension,
                insert_size_sample_size,gap_open,gap_ext,n_hits,prob_cutoff,min_soft_clipped_length,overlap_frac,
                interval_p_cut, output_name,ncores,af,locker,split,ratio,verbose,pid,edit_distance_frac,
                remap_splits,only_discordants,splits,score,insert_size,discordant_filter):
        #I/O
        self.edit_distance_frac = edit_distance_frac
        self.ecc_dna_str = input_bam
        self.qname_bam = qname_bam
        self.sorted_bam_str = sorted_bam
        self.directory = directory
        self.genome_fa = genome_fasta

        #realignment parameters

        # probabilistic realignment options
        self.n_hits = n_hits
        self.prob_cutoff = prob_cutoff
        self.min_sc_length = min_soft_clipped_length
        self.mapq_cutoff = mapq_cutoff
        self.interval_p = interval_p_cut
        self.remap = remap_splits
        self.only_discordants = only_discordants
        self.split = splits
        self.score = score
        self.af= af
        self.insert = insert_size

        # affine gap scoring options
        self.gap_open = gap_open
        self.gap_ext = gap_ext


        #insert size stimation parameters
        self.insert_size_mapq = insert_size_mapq
        self.std_extenstion = std_extension
        self.insert_sample_size = insert_size_sample_size



        #output options

        self.overlap_fraction = overlap_frac
        self.output = output_name
        self.discordant_filter = discordant_filter


        #regular options
        self.cores = ncores
        self.verbose = verbose
        self.lock = locker

        #this two parameters don't work on this class. They are here for printing the parameters
        self.split = split
        self.ratio = ratio

        #for instances running on the same directoiry

        self.pid = pid





    def print_parameters(self):

        print("Running realignment\n")
        print("Probabilistic realignment parameters:\n"
            "\tAlignments to consider: %s \n"
            "\tProbability cut-off to consider as mapped: %s \n"
            "\tMinimum soft-clipped length to attemp realignment: %s \n"
            "\tMinimum bwa mem mapping quality to consider: %s \n"
            "\tGap open penalty: %s \n"
            "\tGap extension penalty: %s \n"
            % (self.n_hits, self.prob_cutoff,self.min_sc_length,self.mapq_cutoff,self.gap_open, self.gap_ext))

        print("Interval extension parameters:\n"
            "\tInsert size mapping quality cut-off: %s \n"
            "\tNumber of read to sample: %s \n"
            "\tNumber of standard deviations to extend the realignment intervals: %s \n"
            % (self.insert_size_mapq,self.insert_sample_size,self.std_extenstion))

        print("eccDNA output options: \n"
            "\tSplit read cut-off: %s \n"
            "\tCoverage ratio cut-off: %s \n" % (self.split,self.ratio))


        print("Interval processing options: \n"
            "\tMerging fraction: %s \n"
            "\tInterval probability cut-off: %s \n"
            % (self.overlap_fraction,self.interval_p))







    def realign(self,peaks):
        """Function that will iterate trough the bam file containing reads indicating eccDNA structural variants and
        will output a bed file containing the soft-clipped reads, the discordant and the coverage within the interval"""

        #open files for every process
        try:
            peaks_pd = pd.DataFrame.from_records(peaks,columns=['chrom', 'start', 'end'])
            sorted_bam = ps.AlignmentFile(self.sorted_bam_str, "rb")
            genome_fa = ps.FastaFile(self.genome_fa)
            ecc_dna = ps.AlignmentFile(self.ecc_dna_str,"rb")

            begin = time.time()








            # compute insert size distribution

            insert_metrics = self.insert


            #define realignment extension interval
            extension = insert_metrics[0] + self.std_extenstion*insert_metrics[1]


            iteration = 0


            results = []
            only_discordants = []



            for index,interval in peaks_pd.iterrows():



                if check_size_and_write(results,only_discordants,self.output,self.lock,self.directory,self.overlap_fraction,self.pid) == True:
                    results = []
                    only_discordants = []

                try:




                    #find out the prior distribution (mate alignment positions).
                    candidate_mates = get_mate_intervals(ecc_dna,interval,self.mapq_cutoff,self.verbose,self.only_discordants)






                    if len(candidate_mates) > 0 or candidate_mates != None:


                        realignment_interval_extended = get_realignment_intervals(candidate_mates,extension,self.interval_p,
                                                                                self.verbose)

                        if realignment_interval_extended is None:
                            continue

                        iteration_results = []
                        iteration_discordants = []
                        disorcordants_per_it = 0
                        for index,mate_interval in realignment_interval_extended.iterrows():

                            iteration += 1

                            #sample realignment intervals
                            #fasta file fetch is 1 based that why I do +1

                            plus_coding_interval = genome_fa.fetch(str(mate_interval['chrom']),int(int(mate_interval['start'])+1),int(int(mate_interval['end'])+1)).upper()
                            interval_length = len(plus_coding_interval)
                            minus_coding_interval = str(Seq(plus_coding_interval).complement())

                            # precompute the denominators of the error model. They will be constants for every interval
                            plus_base_freqs = background_freqs(plus_coding_interval)

                            minus_base_freqs = {'T':plus_base_freqs['A'],'A':plus_base_freqs['T'],
                                                'C':plus_base_freqs['G'],'G':plus_base_freqs['C']}

                            minus_base_freqs = np.array([plus_base_freqs['T'],plus_base_freqs['A'],plus_base_freqs['G'],plus_base_freqs['C']])
                            plus_base_freqs = np.array([plus_base_freqs['A'],plus_base_freqs['T'],plus_base_freqs['C'],plus_base_freqs['G']])


                            #note that I am getting the reads of the interval. Not the reads of the mates

                            for read in ecc_dna.fetch(interval['chrom'],int(interval['start']),int(interval['end']),multiple_iterators=True):


                                if is_soft_clipped(read):

                                    if read.mapq >= self.mapq_cutoff:

                                        # no need to realignment
                                        if read.has_tag('SA') and self.remap != True:


                                            #check realignment from SA tag
                                            support = circle_from_SA(read, self.mapq_cutoff, mate_interval)



                                            if support is  None:
                                                pass

                                            else:

                                                if support['support'] == True:

                                                    score = len(get_longest_soft_clipped_bases(read)['seq'])*  (1-phred_to_prob(np.array(int(read.get_tag('SA').split(',')[4]),dtype=np.float64)))

                                                    #compute mapping positions

                                                    read_end = rightmost_from_read(read)

                                                    supplementary_end = rightmost_from_sa(support['leftmost'],support['cigar'])



                                                    # I store the read name to the output, so that a read counts as 1 no matter it is SC in 2 pieces
                                                    if read.reference_start < support['leftmost']:

                                                        iteration_results.append([interval['chrom'],read.reference_start,(supplementary_end-1),read.qname,iteration,float(round(score,2))])

                                                    elif read.reference_start > support['leftmost']:

                                                        iteration_results.append(
                                                            [interval['chrom'], (support['leftmost']-1), read_end, read.qname,iteration,float(round(score,2))])

                                                    else:
                                                        #uninformative read
                                                        pass



                                        else:
                                            #sc length
                                            sc_len = len(get_longest_soft_clipped_bases(read)['seq'])


                                            if non_colinearity(int(read.cigar[0][0]),int(read.cigar[-1][0]),int(read.pos),
                                                                int(mate_interval.start),int(mate_interval.end)) == True:


                                                if sc_len >= self.min_sc_length:
                                                    edits_allowed = adaptative_myers_k(sc_len, self.edit_distance_frac)
                                                #realignment

                                                    realignment_dict = realign(read,self.n_hits,plus_coding_interval,minus_coding_interval,
                                                                            plus_base_freqs,minus_base_freqs,self.gap_open,self.gap_ext,self.verbose,edits_allowed)


                                                    if realignment_dict == None:

                                                        pass

                                                    else:
                                                        #calc edit distance allowed
                                                        prob = realignment_probability(realignment_dict,interval_length)
                                                        if prob >= self.prob_cutoff and realignment_dict['alignments'][1][3] <= edits_allowed:

                                                            # here I have to retrieve the nucleotide mapping positions. Which should be the
                                                            # the left sampling pysam coordinate - edlib coordinates

                                                            read_end = rightmost_from_read(read)


                                                            soft_clip_start = int(mate_interval['start'])+ int(realignment_dict['alignments'][1][0][0])

                                                            soft_clip_end = int(mate_interval['start']) + int(realignment_dict['alignments'][1][0][1])

                                                            score = sc_len*prob


                                                            # I store the read name to the output, so that a read counts as 1 no matter it is SC in 2 pieces
                                                            if read.reference_start < int(mate_interval['start']) + int(
                                                                    realignment_dict['alignments'][1][0][0]):

                                                                iteration_results.append([interval['chrom'], read.reference_start, soft_clip_end+1, read.qname,iteration,float(round(score,2))])

                                                            elif read.reference_start + int(mate_interval['start']) + int(
                                                                    realignment_dict['alignments'][1][0][0]):

                                                                iteration_results.append([interval['chrom'], soft_clip_start, read_end, read.qname,iteration,float(round(score,2))])

                                                            else:
                                                                # uninformative read
                                                                pass



                                                        else:
                                                            pass
                                            else:

                                                pass

                                    else:
                                        pass
                                else:
                                    #discordant reads
                                    #R2F1 oriented when iterating trough R2
                                    if read.is_reverse == True and read.mate_is_reverse == False:
                                        if read.is_read2:
                                            if read.reference_start < read.next_reference_start:
                                                # discordant read
                                                disorcordants_per_it +=1
                                                iteration_discordants.append([interval['chrom'],read.reference_start,read.next_reference_start + read.infer_query_length(),read.qname])




                                    #R2F1 when iterating trough F1
                                    elif read.is_reverse == False and read.mate_is_reverse ==  True:
                                        if read.is_read2 == False:
                                            if read.next_reference_start < read.reference_start:
                                                disorcordants_per_it +=1
                                                iteration_discordants.append([interval['chrom'], read.next_reference_start,read.reference_start+read.infer_query_length(),                                                                              read.qname])


                        #second pass to add discordant read info
                        if len(iteration_results) > 0:


                            results = results + assign_discordants(iteration_results,iteration_discordants,insert_metrics[0],insert_metrics[1])


                        elif len(iteration_discordants) > 0:
                                discordant_bed = pd.DataFrame.from_records(iteration_discordants,columns=['chrom','start','end','read']).sort_values(['chrom','start','end'])

                                discordant_bed = discordant_bed.groupby(merge_bed(discordant_bed)).agg(
                                    {'chrom': 'first', 'start': 'first', 'end': 'last', 'read': 'count'})


                                for index,disc_interval in discordant_bed.iterrows():
                                    only_discordants.append([disc_interval['chrom'],disc_interval['start'],disc_interval['end'],disc_interval['read'],0])



                except BaseException as e:

                    traceback.print_exc(file=sys.stdout)
                    warnings.warn(
                        "Failed on interval %s due to the error %s" % (
                            str(interval), str(e)))
                    return([1,1])


            ecc_dna.close()
            genome_fa.close()


            # Write process output to disk
            output = iteration_merge(only_discordants,results,
                                    self.overlap_fraction,self.split,self.score,
                                    self.min_sc_length,sorted_bam,self.af,insert_metrics[0],insert_metrics[1],self.discordant_filter)

            write_to_disk(output, self.output, self.lock, self.directory, self.pid)


        except:
            print("Failed on cluster:")
            print(traceback.print_exc(file=sys.stdout))
            return([1,1])

        sorted_bam.close()
        genome_fa.close()
        ecc_dna.close()


        return([0,0])
