from singlecellmultiomics.molecule.molecule import Molecule

import collections

import pandas as pd

class TranscriptMolecule(Molecule):

    def __init__(self, fragment,

             **kwargs):

        self.genes=set()

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

    def _add_fragment(self, fragment):

        self.genes.update(fragment.genes)

        Molecule._add_fragment(self, fragment)

    def write_tags(self):

        for frag in self:

            frag.write_tags()

        Molecule.write_tags(self)

class FeatureAnnotatedMolecule(Molecule):

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

    """

    def __init__(

            self,

            fragment,

            features,

            stranded=None,

            auto_set_intron_exon_features=False,

            capture_locations=False,

            **kwargs):

        """

            Args:

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

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

                stranded : None; not stranded, False: same strand as R1, True: other strand

                capture_locations (bool) : Store information about the locations of the aligned features

                auto_set_intron_exon_features(bool) : obtain intron_exon_features upon initialising

                **kwargs: extra args

        """

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

        self.features = features

        self.hits = collections.defaultdict(set)  # feature -> hit_bases

        self.stranded = stranded

        self.is_annotated = False

        self.capture_locations = capture_locations

        if capture_locations:

            self.feature_locations = {} #feature->locations (chrom,start,end, strand)

        self.junctions = set()

        self.genes = set()

        self.introns = set()

        self.exons = set()

        self.exon_hit_gene_names = set()  # readable names

        self.is_spliced = None

        if auto_set_intron_exon_features:

            self.set_intron_exon_features()

    def set_spliced(self, is_spliced):

        """ Set wether the transcript is spliced, False has priority over True """

        if self.is_spliced and not is_spliced:

            # has already been set

            self.is_spliced = False

        else:

            self.is_spliced = is_spliced

    def set_intron_exon_features(self):

        if not self.is_annotated:

            self.annotate()

        # Collect all hits:

        hits = self.hits.keys()

        # (gene, transcript) -> set( exon_id  .. )

        exon_hits = collections.defaultdict(set)

        intron_hits = collections.defaultdict(set)

        for hit, locations in self.hits.items():

            if not isinstance(hit, tuple):

                continue

            meta = dict(list(hit))

            if 'gene_id' not in meta:

                continue

            if meta.get('type') == 'exon':

                if 'transcript_id' not in meta:

                    continue

                self.genes.add(meta['gene_id'])

                self.exons.add(meta['exon_id'])

                if 'transcript_id' not in meta:

                    raise ValueError(

                        "Please use an Intron GTF file generated using -id 'transcript_id'")

                exon_hits[(meta['gene_id'], meta['transcript_id'])].add(

                    meta['exon_id'])

                if 'gene_name' in meta:

                    self.exon_hit_gene_names.add(meta['gene_name'])

            elif meta.get('type') == 'intron':

                self.genes.add(meta['gene_id'])

                self.introns.add(meta['gene_id'])

        # Find junctions and add all annotations to annotation sets

        debug = []

        for (gene, transcript), exons_overlapping in exon_hits.items():

            # If two exons are detected from the same gene we detected a

            # junction:

            if len(exons_overlapping) > 1:

                self.junctions.add(gene)

                # We found two exons and an intron:

                if gene in self.introns:

                    self.set_spliced(False)

                else:

                    self.set_spliced(True)

            debug.append(

                f'{gene}_{transcript}:' +

                ','.join(

                    list(exons_overlapping)))

        # Write exon dictionary:

        self.set_meta('DB', ';'.join(debug))

    def get_hit_df(self):

        """Obtain dataframe with hits

        Returns:

            pd.DataFrame

        """

        if not self.is_annotated:

            self.set_intron_exon_features()

        d = {}

        tabulated_hits = []

        for hit, locations in self.hits.items():

            if not isinstance(hit, tuple):

                continue

            meta = dict(list(hit))

            for location in locations:

                location_dict = {

                    'chromosome': location[0],

                    'start': location[1][0],

                    'end': location[1][1]}

                location_dict.update(meta)

                tabulated_hits.append(location_dict)

        return pd.DataFrame(tabulated_hits)

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

        # @ todo needs refactor; the psuedoread should just be a Fragment too, solves all issues

        if call_super:

            Molecule.write_tags_to_psuedoreads(self, reads)

        for read in reads:

            if len(self.exons) > 0:

                read.set_tag('EX', ','.join(sorted([str(x) for x in self.exons])))

            else:

                read.set_tag('EX', None)

            if len(self.introns) > 0:

                read.set_tag('IN', ','.join(

                    sorted([str(x) for x in self.introns])))

            else:

                read.set_tag('IN', None)

            if len(self.genes) > 0:

                read.set_tag('GN', ','.join(sorted([str(x) for x in self.genes])))

            else:

                read.set_tag('GN', None)

            if len(self.junctions) > 0:

                read.set_tag('JN', ','.join(

                    sorted([str(x) for x in self.junctions])))

                # Is transcriptome

                read.set_tag('IT', 1)

            elif len(self.genes) > 0:

                # Maps to gene but not junction

                read.set_tag('IT', 0.5)

                read.set_tag('JN', None)

            else:

                # Doesn't map to gene

                read.set_tag('IT', 0)

                read.set_tag('JN', None)

            if self.is_spliced is True:

                read.set_tag('SP', True)

            elif self.is_spliced is False:

                read.set_tag('SP', False)

            if len(self.exon_hit_gene_names):

                read.set_tag('gn', ';'.join(list(self.exon_hit_gene_names)))

            else:

                read.set_tag('gn', None)

    def write_tags(self):

        Molecule.write_tags(self)

        # Write cell ranger tags:

        if self.umi is not None:

            self.set_meta('UB', self.umi)

        bc = list(self.get_barcode_sequences())[0]

        self.set_meta('CB', bc)

        if len(self.exons) > 0:

            self.set_meta('EX', ','.join(sorted([str(x) for x in self.exons])))

        else:

            self.set_meta('EX',None)

        if len(self.introns) > 0:

            self.set_meta('IN', ','.join(

                sorted([str(x) for x in self.introns])))

        else:

            self.set_meta('IN',None)

        if len(self.genes) > 0:

            self.set_meta('GN', ','.join(sorted([str(x) for x in self.genes])))

        else:

            self.set_meta('GN',None)

        if len(self.junctions) > 0:

            self.set_meta('JN', ','.join(

                sorted([str(x) for x in self.junctions])))

            # Is transcriptome

            self.set_meta('IT', 1)

        elif len(self.genes) > 0:

            # Maps to gene but not junction

            self.set_meta('IT', 0.5)

            self.set_meta('JN',None)

        else:

            # Doesn't map to gene

            self.set_meta('IT', 0)

            self.set_meta('JN', None)

        if self.is_spliced is True:

            self.set_meta('SP', True)

        elif self.is_spliced is False:

            self.set_meta('SP', False)

        if len(self.exon_hit_gene_names):

            self.set_meta('gn', ';'.join(list(self.exon_hit_gene_names)))

        else:

            self.set_meta('gn',None)

    def annotate(self, method=0):

        """

            Args:

                method (int) : 0, obtain blocks and then obtain features. 1, try to obtain features for every aligned base

        """

        # When self.stranded is None, set to None strand. If self.stranded is

        # True reverse the strand, otherwise copy the strand

        strand = None if self.stranded is None else '+-'[

            (not self.strand if self.stranded else self.strand)]

        self.is_annotated = True

        if method == 0:

            # Obtain all blocks:

            try:

                for start, end in self.get_aligned_blocks():

                    for hit in self.features.findFeaturesBetween(

                            chromosome=self.chromosome, sampleStart=start, sampleEnd=end, strand=strand):

                        hit_start, hit_end, hit_id, hit_strand, hit_ids = hit

                        self.hits[hit_ids].add(

                            (self.chromosome, (hit_start, hit_end)))

                        if self.capture_locations:

                            if not hit_id in self.feature_locations:

                                self.feature_locations[hit_id] = []

                            self.feature_locations[hit_id].append( (hit_start, hit_end, hit_strand))

            except TypeError:

                # This happens when no reads map

                pass

        else:

            for read in self.iter_reads():

                for q_pos, ref_pos in read.get_aligned_pairs(

                        matches_only=True, with_seq=False):

                    for hit in self.features.findFeaturesAt(

                            chromosome=read.reference_name, lookupCoordinate=ref_pos, strand=strand):

                        hit_start, hit_end, hit_id, hit_strand, hit_ids = hit

                        self.hits[hit_ids].add((read.reference_name, ref_pos))

                        if self.capture_locations:

                            if not hit_id in self.feature_locations:

                                self.feature_locations[hit_id] = []

                            self.feature_locations[hit_id].append( (hit_start, hit_end, hit_strand))