from singlecellmultiomics.fragment import Fragment
from singlecellmultiomics.utils.sequtils import hamming_distance


class CHICFragment(Fragment):
    def __init__(self,
                 reads,
                 R1_primer_length=4,
                 R2_primer_length=6, # Is automatically detected now
                 assignment_radius=0,
                 umi_hamming_distance=1,
                 invert_strand=False,
                 no_umi_cigar_processing=False,
                 **kwargs
                 ):
        self.invert_strand = invert_strand


        # Perform random primer autodect,
        # When the demux profile is MX:Z:scCHIC384C8U3l, then there is no random primer
        for read in reads:
            if read is not None and read.has_tag('MX'):
                if read.get_tag('MX')[-1]=='l':
                    R2_primer_length = 0
                break

        Fragment.__init__(self,
                          reads,
                          assignment_radius=assignment_radius,
                          R1_primer_length=R1_primer_length,
                          R2_primer_length=R2_primer_length,
                          umi_hamming_distance=umi_hamming_distance,
                          max_NUC_stretch = 18,
                          **kwargs

                )

        # set CHIC cut site given reads
        self.no_umi_cigar_processing = no_umi_cigar_processing
        self.strand = None
        self.ligation_motif = None
        self.site_location = None
        self.cut_site_strand = None
        self.found_valid_site = False
        self.identify_site()


        if self.is_valid():

            if self.assignment_radius==0:
                self.match_hash = (
                    self.strand,
                    self.cut_site_strand,
                    self.site_location[0],
                    self.site_location[1],
                    self.sample)
            else:
                self.match_hash = (
                    self.cut_site_strand,
                    self.site_location[0],
                    self.sample)
        else:
            self.match_hash = None

    def set_site(
            self,
            site_chrom,
            site_pos,
            site_strand=None,
            is_trimmed=False):
        if self.found_valid_site:
            self.set_meta('DS', site_pos)
        if site_strand is not None:
            self.set_meta('RS', site_strand)
        self.set_strand(site_strand)
        self.site_location = (site_chrom, site_pos)
        self.cut_site_strand = site_strand

    def identify_site(self):
        self.found_valid_site = False
        R1 = self.get_R1()

        if R1 is None:
            self.set_rejection_reason("R1_undefined")
            return None

        if R1.has_tag('lh'):
            self.ligation_motif = R1.get_tag('lh')

        """ Valid configs:
        Observed read pair:
        T######## R1 ########## ^ ########## R2 ##########
        Real molecule:
        A###########################################################

        Real molecule: and adapter:
        ADAPTER: 5' NNNNT 3'
                        A{A:80%,N:20%}NNN [CHIC MOLECULE]
                                      ^ real cut site
        """

        is_trimmed = (R1.has_tag('MX') and R1.get_tag(
            'MX').startswith('scCHIC'))

        if R1.is_unmapped:
            self.set_rejection_reason("R1_unmapped")
            return(None)

        ### Check if the orientation of the reads makes sense, reads need to point inwards
        # | ----- >   < ---- |
        if self.has_R2():
            R2 = self.get_R2()
            if not R2.is_unmapped:
                if R1.is_reverse and not R2.is_reverse:
                    pass
                elif not R1.is_reverse and R2.is_reverse:
                    pass
                else:
                    # it does not make sense ...
                    self.set_rejection_reason("orientation")
                    return(None)


        # Identify the start coordinate of Read 1 by reading the amount of softclips on the start of the read
        r1_start =(R1.reference_end+1 if R1.is_reverse else R1.reference_start-2)
        if not self.no_umi_cigar_processing:
            if R1.is_reverse:
                if R1.cigartuples[-1][0]==4: # softclipped at end
                    r1_start+=R1.cigartuples[-1][1]
            else:
                if R1.cigartuples[0][0]==4: # softclipped at start
                    r1_start-=R1.cigartuples[0][1]

        if is_trimmed:
            # The first base of the read has been taken off and the lh tag is
            # already set, this can be copied to RZ
            self.found_valid_site = True
            self.set_site(
                site_strand=not R1.is_reverse if self.invert_strand else R1.is_reverse,
                # We sequence the other strand (Starting with a T, this is an A in the molecule), the digestion thus happened on the other strand
                # On the next line we asume that the mnsase cut is one base after the ligated A, but it can be more bases upstream
                site_chrom=R1.reference_name,
                site_pos=r1_start,
                is_trimmed=True
            )

        else:
            self.found_valid_site = True
            self.set_site(
                site_strand=not R1.is_reverse if self.invert_strand else R1.is_reverse,
                # We sequence the other strand (Starting with a T, this is an A in the molecule), the digestion thus happened on the other strand
                # On the next line we asume that the mnsase cut is one base after the ligated A, but it can be more bases upstream
                site_chrom=R1.reference_name,
                site_pos=(r1_start - 1 if R1.is_reverse else r1_start + 1),
                is_trimmed=False)

    def is_valid(self):
        if self.qcfail:
            return False

        if self.max_fragment_size is not None:
            try:
                size = self.get_fragment_size()
                if size>self.max_fragment_size:
                    return False
            except Exception:
                pass

        return self.found_valid_site

    def get_site_location(self):
        if self.site_location is not None:
            return self.site_location
        else:
            # We need some kind of coordinate...
            for read in self:
                if read is not None and read.reference_name is not None and read.reference_start is not None:
                    return read.reference_name, read.reference_start

    def __repr__(self):
        return Fragment.__repr__(
            self) + f'\n\tMNase cut site:{self.get_site_location()}'

    def __eq__(self, other):
        # Make sure fragments map to the same strand, cheap comparisons
        if self.match_hash != other.match_hash:
            return False

        if self.site_location is None or other.site_location is None:
            return False

        # Check distance between the fragments:
        if self.assignment_radius>0 and abs(self.site_location[1]-other.site_location[1])>self.assignment_radius:
            return False


        # Make sure UMI's are similar enough, more expensive hamming distance
        # calculation
        return self.umi_eq(other)