#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 circlemap.utils import *
import pandas as pd
import traceback
import multiprocessing as mp
import warnings
import datetime



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

    queue = mp.Manager().Queue()

    def __init__(self, input_bam,output,qname_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,
                interval_p_cut,ncores,locker,verbose,pid,edit_distance_frac,
                remap_splits,only_discordants,score,insert_size,manager):
        #I/O
        self.edit_distance_frac = edit_distance_frac
        self.ecc_dna_str = input_bam
        self.output = output
        self.qname_bam = qname_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.score = score
        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

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



        #for instances running on the same directoiry

        self.pid = pid

        #parallel enviroment
        self.read_list = manager.list()
        self.read_count = manager.Value('i', 0)
        self.write_round = manager.Value('i', 0)







    def listener_writer(self,bam):

        f = open('test.sam',"w")

        header = bam.header

        while True:

            # Read from the queue and do nothing
            read = self.queue.get()


            if read == "DONE":
                f.close()
                print("breaking")
                bam.close()
                break
            else:

                pysam_read = ps.AlignedSegment.fromstring(read,bam.header)
                f.write(read + "\n")
                bam.write(pysam_read)

    def kill(self):
        print("KILLING")
        self.queue.put("DONE")

    def beta_version_warning(self):
        """Warn the user that this is experimental"""
        print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S: You are using a beta version feature"))
        warnings.warn("The bam2bam feature on Circle-Map is experimental. The development of this feature is active, but"
                    " have in mind that it might produce unintended results. Check https://github.com/iprada/Circle-Map"
                    " for the development status.")




    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'])
            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



            for index,interval in peaks_pd.iterrows():




                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 = []
                        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:
                                                    self.queue.put(read.to_string())

                                                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)

                                                            #aln start on the reference
                                                            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
                                                            # Soft-clipped aligned upstream. Primary aligned downstream
                                                            if read.reference_start < int(mate_interval['start']) + int(
                                                                    realignment_dict['alignments'][1][0][0]):
                                                                    # construct tag
                                                                    sa_tag = realignment_read_to_SA_string(realignment_dict,
                                                                                                prob, interval['chrom'],
                                                                                                soft_clip_start)


                                                                    #read.tags += [('SA', sa_tag)]

                                                                    self.queue.put(read.to_string())



                                                            # soft-clipped aligned downstream primary alignment is upstream
                                                            elif read.reference_start + int(mate_interval['start']) + int(
                                                                    realignment_dict['alignments'][1][0][0]):

                                                                sa_tag = realignment_read_to_SA_string(realignment_dict,
                                                                                                    prob, interval[
                                                                                                        'chrom'],
                                                                                                    soft_clip_start)

                                                                read.tags += [('SA', sa_tag)]

                                                                self.queue.put(read.to_string())

                                                            else:
                                                                # uninformative read
                                                                pass



                                                        else:
                                                            pass
                                            else:

                                                pass

                                    else:
                                        pass
                                else:
                                    pass


                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()

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

        genome_fa.close()
        ecc_dna.close()


        return([0,0])
