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

class ScarTraceFragment(Fragment):
    """
    Fragment definition for ScarTrace
    """

    def __init__(self, reads,scartrace_r1_primers=None, **kwargs):
        Fragment.__init__(self, reads,  **kwargs)
        self.scartrace_r1_primers = scartrace_r1_primers
        assert scartrace_r1_primers is not None, 'please supply primer sequences'

    # remove the set_umi function
    def set_umi(self, **kwargs):
        pass

    def is_valid(self):
        """ Check if R1 starts with the defined primers and if R2 and R1 are mapped

        Returns:
            bool
        """
        if not self.has_R1() or not self.has_R2():
            if not self.has_R1():
                self.set_meta('fr','no_R1', as_set=True)
            if not self.has_R2():
                self.set_meta('fr','no_R2', as_set=True)
            return False

        if self[0].is_unmapped or self[1].is_unmapped:
            self.set_meta('fr','unmapped_mate', as_set=True)
            return False

        r1_seq = self.get_R1().seq

        if self.get_R1().is_reverse:
            r1_seq = reverse_complement(r1_seq)

        if any( r1_seq.startswith(primer) for primer in self.scartrace_r1_primers ):
            # Good
            return True

        self.set_meta('fr','primer_not_matching', as_set=True)
        # Bad
        return False




    # Replace the equality function
    def __eq__(self, other):  # other can also be a Molecule!
        # Make sure fragments map to the same strand, cheap comparisons
        if self.sample != other.sample:
            return False

        if self.strand != other.strand:
            return False

        if min(abs(self.span[1] -
                   other.span[1]), abs(self.span[2] -
                                       other.span[2])) > self.assignment_radius:
            return False

        # Sample matches and starting position is within the defined span
        # radius
        return True