from singlecellmultiomics.molecule.molecule import Molecule

from singlecellmultiomics.molecule.featureannotatedmolecule import FeatureAnnotatedMolecule

class NlaIIIMolecule(Molecule):

    """Nla III based Molecule class

    Args:

        fragments (singlecellmultiomics.fragment.Fragment): Fragments to associate to the molecule

        **kwargs: extra args

    """

    def __init__(self, fragment,

                 site_has_to_be_mapped=False,

                 # molecule is invalid when NLA site does not map  (requires

                 # reference)

                 **kwargs):

        self.site_location = None

        self.site_has_to_be_mapped = site_has_to_be_mapped

        Molecule.__init__(self, fragment, **kwargs)

    def _add_fragment(self,fragment):

        # Update the cut coordinate tho the (left most extreme value)

        self.assignment_radius = fragment.assignment_radius

        if fragment.site_location is not None:

            if self.site_location is None:

                self.site_location = [fragment.site_location[0],fragment.site_location[1]]

            elif fragment.strand: # Reverse:

                self.site_location[1] = max(fragment.site_location[1], self.site_location[1]) # this is the coordinate

            else:

                self.site_location[1] = min(fragment.site_location[1], self.site_location[1]) # this is the coordinate

        # else : writing a fragment which has no cut location associated

        Molecule._add_fragment(self, fragment)

    def write_tags(self):

        Molecule.write_tags(self)

        if self.reference is not None:  # Only calculate this statistic when a reference is available

            if self.get_cut_site() is not None:

                self.set_meta('Us', self.get_undigested_site_count())

    def is_valid(self, set_rejection_reasons=False, reference=None):

        try:

            chrom, start, strand = self.get_cut_site()

        except Exception as e:

            if set_rejection_reasons:

                self.set_rejection_reason('no_cut_site_found')

            return False

        if self.site_has_to_be_mapped:

            if reference is None:

                if self.reference is None:

                    raise ValueError(

                        'Please supply a reference (PySAM.FastaFile)')

            reference = self.reference

            if reference.fetch(chrom, start, start + 4).upper() != 'CATG':

                if set_rejection_reasons:

                    self.set_rejection_reason('no_CATG_in_ref')

                return False

        return True

    def get_fragment_span_sequence(self, reference=None):

        """Obtain the sequence between the start and end of the molecule

        Args:

            reference(pysam.FastaFile) : reference  to use.

                If not specified `self.reference` is used

        Returns:

            sequence (str)

        """

        if self.chromosome is None:

            return ''

        if reference is None:

            if self.reference is None:

                raise ValueError('Please supply a reference (PySAM.FastaFile)')

        reference = self.reference

        return reference.fetch(

            self.chromosome,

            self.spanStart,

            self.spanEnd).upper()

    def get_undigested_site_count(self, reference=None):

        """

        Obtain the amount of undigested sites in the span of the molecule

        Args:

            reference(pysam.FastaFile) : reference handle

        Returns:

            undigested_site_count : int

            amount of undigested cut sites in the mapping span of the molecule

        Raises:

            ValueError : when the span of the molecule is not properly defined

        """

        if reference is None:

            if self.reference is None:

                raise ValueError('Please supply a reference (PySAM.FastaFile)')

        reference = self.reference

        seq = self.get_fragment_span_sequence(reference)

        total = seq.count('CATG')

        if self.strand == 0 and seq.endswith('CATG'):

            total -= 1

        elif self.strand == 1 and seq.startswith('CATG'):

            total -= 1

        return total

    def write_tags_to_psuedoreads(self, reads, call_super=True):

        if call_super:

            Molecule.write_tags_to_psuedoreads(self, reads)

        if self.reference is not None:  # Only calculate this statistic when a reference is available

            if self.get_cut_site() is not None:

                us = self.get_undigested_site_count()

                for read in reads:

                    read.set_tag('Us', us)

class AnnotatedNLAIIIMolecule(FeatureAnnotatedMolecule, NlaIIIMolecule):

    """Nla III based Molecule which is annotated with features (genes/exons/introns, .. )

    Args:

        fragments (singlecellmultiomics.fragment.Fragment): Fragments to associate to the molecule

        features (singlecellmultiomics.features.FeatureContainer) : container to use to obtain features from

        **kwargs: extra args

    """

    def write_tags(self):

        NlaIIIMolecule.write_tags(self)

        FeatureAnnotatedMolecule.write_tags(self)

    def __init__(self, fragment, features, **kwargs):

        self.site_location = None

        FeatureAnnotatedMolecule.__init__(self, fragment, features, **kwargs)

        NlaIIIMolecule.__init__(self, fragment, **kwargs)

    def write_tags_to_psuedoreads(self,reads):

        NlaIIIMolecule.write_tags_to_psuedoreads(self, reads)

        FeatureAnnotatedMolecule.write_tags_to_psuedoreads(self,reads,call_super=False)