from singlecellmultiomics.molecule.molecule import Molecule
from singlecellmultiomics.molecule.featureannotatedmolecule import FeatureAnnotatedMolecule
import collections

class CHICMolecule(Molecule):
    """CHIC Molecule class

    Args:
        fragments (singlecellmultiomics.fragment.Fragment): Fragments to associate to the molecule
        **kwargs: extra args

    """

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

        self.site_location = None
        Molecule.__init__(self, fragment, **kwargs)
        self.ligation_motif = None

        # Extracted from fragments:
        self.assignment_radius=None



    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 get_cut_site(self):
        """For restriction based protocol data, obtain genomic location of cut site

        Returns:
            None if site is not available

            chromosome (str)
            position (int)
            strand (bool)
        """
        return (*self.site_location, self.strand)

    def write_tags(self):

        # Write DS tag when it could have been updated
        if self.assignment_radius is not None and self.assignment_radius>0 and len(self)>1:
            for frag in self:
                frag.set_meta('DS', self.site_location[1])

        Molecule.write_tags(self)

    def is_valid(self, set_rejection_reasons=False):

        try:
            chrom, start, strand = self.get_cut_site()
        except Exception as e:
            if set_rejection_reasons:
                self.set_rejection_reason(f'no_cut_site_found')
            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 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 __finalise__(self):
        # Obtain ligation motif(s)
        self.update_ligation_motif()
        Molecule.__finalise__(self)

    def update_ligation_motif(self):
        """
        Extract lh tag from associated reads and set most common one to the RZ tag
        """
        ligation_motif_counts = collections.Counter()
        for fragment in self:
            if fragment.ligation_motif is not None:
                ligation_motif_counts[fragment.ligation_motif]+=1
        if len(ligation_motif_counts)>0:
            self.ligation_motif = ligation_motif_counts.most_common(1)[0][0]
            self.set_meta('RZ', self.ligation_motif)


class CHICNLAMolecule(Molecule):
    """CHIC NLA Molecule class

    Args:
        fragments (singlecellmultiomics.fragment.Fragment): Fragments to associate to the molecule
        **kwargs: extra args

    """
    def get_upstream_site(self, scan_extra_bp=3):
        read = None
        for fragment in self:
            if fragment[0] is not None:
                read = fragment[0]
                break

        if read is None:
            raise ValueError() # Read 1 is not mapped

        if self.strand==0:
            start = read.reference_start
            return self.reference.fetch( self.chromosome, start-4-scan_extra_bp, start)
        else:
            start = read.reference_end
            return self.reference.fetch( self.chromosome, start, start+4+scan_extra_bp)


    def __init__(self, fragment, reference, **kwargs):
        self.reference = reference
        CHICMolecule.__init__(self, fragment, reference=reference, **kwargs)

    def write_tags(self):
        try:
            site = self.get_upstream_site()
            self.set_meta('RZ',site)
            if 'N' in site:
                self.set_meta('dt','UNK')
            elif 'CATG' in site:
                self.set_meta('dt','NLA')
            else:
                self.set_meta('dt','CHIC')
        except ValueError:
            self.set_meta('dt','UNK')
            self.set_meta('RR','No_R1')




class AnnotatedCHICMolecule(CHICMolecule, FeatureAnnotatedMolecule):
    """Chic 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):
        CHICMolecule.write_tags(self)
        FeatureAnnotatedMolecule.write_tags(self)

    def __init__(self, fragment, features, **kwargs):
        CHICMolecule.__init__(self, fragment, **kwargs)
        FeatureAnnotatedMolecule.__init__(self, fragment, features, **kwargs)