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

import pysam as ps
import os
from utils import *
import time
import sys
import warnings


class readExtractor:
    """Class for managing the read extracting part of circle map"""
    def __init__(self,sorted_bam,output_bam,working_dir,mapq_cutoff,extract_discordant,extract_soft_clipped,extract_hard_clipped,
                verbose,parser
                ):
        #input-output
        self.sorted_bam = sorted_bam
        self.output_bam = output_bam
        #working place
        self.working_dir = working_dir

        #read options
        self.no_discordants = extract_discordant
        self.no_soft_clipped = extract_soft_clipped
        self.no_hard_clipped = extract_hard_clipped

        #mapq cutoff

        self.mapq_cutoff = mapq_cutoff

        #verbose level
        self.verbose = int(verbose)
        #parser options
        self.parser = parser

    def extract_sv_circleReads(self):

        """Function that extracts Structural Variant reads that indicate circular DNA,
        The programme with extract soft-clipped reads and R2F1 (<- ->) oriented reads"""

        os.chdir(self.working_dir)

        #input
        raw_bam = ps.AlignmentFile(self.working_dir + "/" + self.sorted_bam, "rb")

        #HD the tag for the header line. SO indicates sorting order of the alignements
        if 'HD' in raw_bam.header:

            if raw_bam.header['HD']['SO'] != 'queryname':
                sys.stderr.write(
                    "The input bam header says that bam is not sorted by queryname. It is sorted by %s\n\n" % (raw_bam.header['HD']['SO']))
                sys.stderr.write(
                    "Sort your bam file queryname with the following command:\n\n\tsamtools sort -n -o output.bam input.bam")

                time.sleep(0.01)

                self.parser.print_help()
                sys.exit(1)
        else:

            if self.verbose >=2:
                warnings.warn("WARNING:Circle-Map does not know if the input bam is queryname sorted\n Please check that, the output would be unexpected otherwise")
                print("As sanity check, sort your bam file queryname with the following command:\n\n\tsamtools sort -n -o output.bam input.bam")




        circle_sv_reads = ps.AlignmentFile(self.working_dir + "/" + self.output_bam , "wb", template=raw_bam)


        #modify the tag to unsorted
        if 'HD' in raw_bam.header == True:
            circle_sv_reads.header['HD']['SO'] = 'unsorted'

        if self.verbose >=3:
            print("Extracting circular structural variants")

        #timing
        begin = time.time()


        #cache read1. operate in read2. this speed-ups the search
        read1 = ''

        #counter for processed reads
        processed_reads = 0

        for read in raw_bam:

            if self.verbose >=3:
                processed_reads +=1


                if (processed_reads/1000000).is_integer() == True:
                    partial_timer = time.time()
                    partial_time = (partial_timer - begin)/60
                    print("Processed %s reads in %s mins" % (processed_reads,round(partial_time,3)))

            if read.is_read1:
                read1 = read
            # Checks for initialization read1 = '' by looking for read1 type == string.
            # read1 type == pysam.Alignment after first read.is_read1 is encountered
            # Condition should only be met if unpaired read2 begins BAM file
            elif isinstance(read1, str):
                pass
            else:
                if read.is_read2 and read.qname == read1.qname:
                    # both reads in memory
                    read2 = read

                    #both reads need to be mapped
                    if read1.is_unmapped == False and read2.is_unmapped == False:


                        if read2.is_reverse and read1.is_reverse == False:


                            # read2 leftmost mapping position smaller than read1 leftmost mapping position
                            if read2.reference_start < read1.reference_start:


                                #aligned to the same chromosome
                                if read1.reference_id == read2.reference_id:


                                    if read1.mapq >= self.mapq_cutoff and read2.mapq >= self.mapq_cutoff:


                                        #is discordant extraction turn off?

                                        if self.no_discordants == False:

                                            #add mate mapping quality info
                                            read1.tags += [('MQ',read2.mapq)]
                                            read2.tags += [('MQ', read1.mapq)]

                                            circle_sv_reads.write(read1)
                                            circle_sv_reads.write(read2)
                                        else:

                                            pass
                                    else:

                                        #extract soft-clipped if the mapq is high enough
                                        write_clipped_read(circle_sv_reads, read1, read2, self.no_soft_clipped,
                                                        self.no_hard_clipped, self.mapq_cutoff)

                                        write_clipped_read(circle_sv_reads, read2, read1, self.no_soft_clipped,
                                                        self.no_hard_clipped, self.mapq_cutoff)



                                else:

                                    write_clipped_read(circle_sv_reads, read1, read2, self.no_soft_clipped,
                                                    self.no_hard_clipped, self.mapq_cutoff)

                                    write_clipped_read(circle_sv_reads, read2, read1, self.no_soft_clipped,
                                                    self.no_hard_clipped, self.mapq_cutoff)


                            else:

                                #if the leftmost mapping condition is not met check if they are soft-clipped
                                write_clipped_read(circle_sv_reads, read1, read2, self.no_soft_clipped,
                                                self.no_hard_clipped, self.mapq_cutoff)

                                write_clipped_read(circle_sv_reads, read2, read1, self.no_soft_clipped,
                                            self.no_hard_clipped, self.mapq_cutoff)


                        else:

                            #check soft-clipped if R2F1 orientation is not True

                                write_clipped_read(circle_sv_reads, read1, read2, self.no_soft_clipped,
                                                    self.no_hard_clipped, self.mapq_cutoff)

                                write_clipped_read(circle_sv_reads, read2, read1, self.no_soft_clipped,
                                            self.no_hard_clipped, self.mapq_cutoff)


                    else:

                        #check read 1 and read two for independent unmaps
                        if read1.is_unmapped == False:
                            write_clipped_read(circle_sv_reads, read1,read2, self.no_soft_clipped,
                                                self.no_hard_clipped, self.mapq_cutoff, own_mapq=True)

                        if read2.is_unmapped == False:
                            write_clipped_read(circle_sv_reads, read2, read1, self.no_soft_clipped,
                                                self.no_hard_clipped, self.mapq_cutoff, own_mapq=True)

                else:
                    # reads are not queryname sorted and cannot be processed in paired mode
                    warnings.warn("Unpaired reads found. Is your bam file queryname sorted?")


        end = time.time()

        circle_sv_reads.close()



        if self.verbose >=3:


            print("finished extracting reads. Elapsed time:", (end - begin) / 60, "mins")

            print("Thanks for using Circle-Map")
