from singlecellmultiomics.utils.sequtils import hamming_distance

import pysamiterators.iterators

import singlecellmultiomics.bamProcessing

from singlecellmultiomics.fragment import Fragment

from array import array

import itertools

import numpy as np

from singlecellmultiomics.utils import style_str, prob_to_phred, phredscores_to_base_call, base_probabilities_to_likelihood, likelihood_to_prob

import textwrap

import singlecellmultiomics.alleleTools

import functools

import typing

import pysam

import pysamiterators

from singlecellmultiomics.utils import find_ranges, create_MD_tag

import pandas as pd

from uuid import uuid4

from cached_property import cached_property

from collections import Counter, defaultdict

###############

# Variant validation function

def detect_alleles(molecules,

                   contig,

                   position,

                   min_cell_obs=3,

                   base_confidence_threshold=None,

                   classifier=None):  # [ alleles]

    """

    Detect the alleles (variable bases) present at the selected location

    Args:

        molecules : generator to extract molecules from

        variant_location(tuple) : (contig, position) zero based location of the location to test

        min_cell_obs (int) : minimum amount of cells containing the allele to be emitted

        confidence_threshold(float) : minimum confidence of concensus base-call to be taken into account

        classifier (obj) : classifier used for consensus call, when no classifier is supplied a mayority vote is used

    """

    observed_alleles = defaultdict(set)  # cell -> { base_call , .. }

    for molecule in molecules:

        base_call = molecule.get_consensus_base(contig, position, classifier=classifier)

        # confidence = molecule.get_mean_base_quality(*variant_location, base_call)

        if base_call is not None:

            observed_alleles[base_call].add(molecule.sample)

    return [allele for allele, cells in observed_alleles.items() if len(cells) >= min_cell_obs]

def get_variant_phase(molecules, contig, position, variant_base, allele_resolver,

                      phasing_ratio_threshold=None):  # (location,base) -> [( location, base, idenfifier)]

    alleles = [variant_base]

    phases = defaultdict(Counter)  # Allele_id -> variant->obs

    for molecule in molecules:

        # allele_obs = molecule.get_allele(return_allele_informative_base_dict=True,allele_resolver=allele_resolver)

        allele = list(molecule.get_allele(allele_resolver))

        if allele is None or len(allele) > 1 or len(allele) == 0:

            continue

        allele = allele[0]

        base = molecule.get_consensus_base(contig, position)

        if base in alleles:

            phases[base][allele] += 1

        else:

            pass

    if len(phases[variant_base]) == 0:

        raise ValueError("Phasing not established, no gSNVS available")

    phased_allele_id = phases[variant_base].most_common(1)[0][0]

    # Check if the phasing noise is below the threshold:

    if phasing_ratio_threshold is not None:

        correct = phases[variant_base].most_common(1)[0][1]

        total = sum(phases[variant_base].values())

        phasing_ratio = correct / total

        if correct / total < phasing_ratio_threshold:

            raise ValueError(f'Phasing ratio not met. ({phasing_ratio}) < {phasing_ratio_threshold}')

    # Check if the other allele i

    return phased_allele_id

###############

@functools.lru_cache(maxsize=1000)

def might_be_variant(chrom, pos, variants, ref_base=None):

    """Returns True if a variant exists at the given coordinate"""

    if ref_base == 'N':

        return False

    try:

        for record in variants.fetch(chrom, pos, pos + 1):

            return True

    except ValueError as e:

        return False # Happens when the contig does not exists, return False

    return False

def consensii_default_vector():

    """a numpy vector with 5 elements initialsed as zeros"""

    return np.zeros(5)

# 1: read1 2: read2

def molecule_to_random_primer_dict(

        molecule,

        primer_length=6,

        primer_read=2,

        max_N_distance=0):

    rp = defaultdict(list)

    # First add all reactions without a N in the sequence:

    for fragment in molecule:

        h_contig, hstart, hseq = fragment.get_random_primer_hash()

        if hseq is None:

            # This should really not happen with freshly demultiplexed data, it means we cannot extract the random primer sequence

            # which should be present as a tag (rS) in the record

            rp[None, None, None].append(fragment)

        elif 'N' not in hseq:

            rp[h_contig, hstart, hseq].append(fragment)

    # Try to match reactions with N with reactions without a N

    for fragment in molecule:

        h_contig, hstart, hseq = fragment.get_random_primer_hash()

        if hseq is not None and 'N' in hseq:

            # find nearest

            for other_contig, other_start, other_seq in rp:

                if other_contig != h_contig or other_start != hstart:

                    continue

                if hseq.count('N') > max_N_distance:

                    continue

                if 'N' in other_seq:

                    continue

                if hamming_distance(hseq, other_seq) == 0:

                    rp[other_contig, other_start, other_seq].append(fragment)

    return rp

class Molecule():

    """Molecule class, contains one or more associated fragments

    Attributes:

        fragments (list): associated fragments

        spanStart (int): starting coordinate of molecule, None if not available

        spanEnd (int): ending coordinate of molecule, None if not available

        chromosome (str): mapping chromosome of molecule, None if not available

        cache_size (int): radius of molecule assignment cache

        reference (pysam.FastaFile) : reference file, used to obtain base contexts

            and correct aligned_pairs iteration when the MD tag is not correct

        strand (bool): mapping strand.

            True when strand is REVERSE

            False when strand is FORWARD

            None when strand is not determined

    """

    def get_empty_clone(self, fragments=None):

        return type(self)(fragments,

                          cache_size = self.cache_size,

                          reference=self.reference,

                          min_max_mapping_quality=self.min_max_mapping_quality,

                          allele_assingment_method=self.allele_assingment_method,

                          allele_resolver=self.allele_resolver,

                          mapability_reader=self.mapability_reader,

                          max_associated_fragments=self.max_associated_fragments,

                          **self.kwargs)

    def __init__(self,

                 fragments: typing.Optional[typing.Iterable] = None,

                 cache_size: int = 10_000,

                 reference: typing.Union[pysam.FastaFile, pysamiterators.CachedFasta] = None,

                 # When all fragments have a mapping quality below this value

                 # the is_valid method will return False,

                 min_max_mapping_quality: typing.Optional[int] = None,

                 mapability_reader: typing.Optional[singlecellmultiomics.bamProcessing.MapabilityReader] = None,

                 allele_resolver: typing.Optional[singlecellmultiomics.alleleTools.AlleleResolver] = None,

                 max_associated_fragments=None,

                 allele_assingment_method=1, # 0: all variants from the same allele, 1: likelihood

                 **kwargs

                 ):

        """Initialise Molecule

        Parameters

        ----------

        fragments :  list(singlecellmultiomics.fragment.Fragment)

            Fragment to assign to Molecule. More fragments can be added later

        min_max_mapping_quality :  When all fragments have a mapping quality below this value the is_valid method will return False

        allele_resolver :  alleleTools.AlleleResolver or None. Supply an allele resolver in order to assign an allele to the molecule

        mapability_reader : singlecellmultiomics.bamProcessing.MapabilityReader, supply a mapability_reader to set mapping_quality of 0 to molecules mapping to locations which are not mapping uniquely during in-silico library generation.

        cache_size (int): radius of molecule assignment cache

        max_associated_fragments(int) : Maximum amount of fragments associated to molecule. If more fragments are added using add_fragment() they are not added anymore to the molecule

        """

        self.kwargs = kwargs

        self.reference = reference

        self.fragments = []

        self.spanStart = None

        self.spanEnd = None

        self.chromosome = None

        self.cache_size = cache_size

        self.strand = None

        self.umi = None

        self.overflow_fragments = 0

        self.umi_hamming_distance = None

        # when set, when comparing to a fragment the fragment to be added has

        # to match this hash

        self.fragment_match = None

        self.min_max_mapping_quality = min_max_mapping_quality

        self.umi_counter = Counter()  # Observations of umis

        self.max_associated_fragments = max_associated_fragments

        if fragments is not None:

            if isinstance(fragments, list):

                for frag in fragments:

                    self.add_fragment(frag)

            else:

                self.add_fragment(fragments)

        self.allele_resolver = allele_resolver

        self.mapability_reader = mapability_reader

        self.allele_assingment_method = allele_assingment_method

        self.methylation_call_dict = None

        self.finalised = False

        self.obtained_allele_likelihoods = None

    @cached_property

    def can_be_split_into_allele_molecules(self):

        l = self.allele_likelihoods

        if l is None or len(l)<=1:

            return False

        return True

    def split_into_allele_molecules(self):

        """

        Split this molecule into multiple molecules, associated to multiple alleles

        Returns:

            list_of_molecules: list

        """

        # Perform allele based clustering

        allele_clustered_frags = {}

        for fragment in self:

            n = self.get_empty_clone(fragment)

            if n.allele not in allele_clustered_frags:

                allele_clustered_frags[n.allele] = []

            allele_clustered_frags[n.allele].append(n)

        allele_clustered = {}

        for allele, assigned_molecules in allele_clustered_frags.items():

            for i, m in enumerate(assigned_molecules):

                if i == 0:

                    allele_clustered[allele] = m

                else:

                    allele_clustered[allele].add_molecule(m)

        if len(allele_clustered)>1:

            for m in allele_clustered.values():

                m.set_meta('cr', 'SplitUponAlleleClustering')

        return list(allele_clustered.values())

    @cached_property

    def allele(self):

        if self.allele_resolver is None:

            return None

        if self.allele_assingment_method == 0:

            # Obtain allele if available

            if self.allele_resolver is not None:

                try:

                    hits = self.get_allele(self.allele_resolver)

                    # Only store when we have a unique single hit:

                    if len(hits) == 1:

                        self.allele = list(hits)[0]

                except ValueError as e:

                    # This happens when a consensus can not be obtained

                    pass

        elif self.allele_assingment_method == 1:

            al = Counter(self.allele_likelihoods)

            if al is None or len(al)<1:

                return None

            return al.most_common(1)[0][0]

        raise NotImplementedError(f'allele_assingment_method {self.allele_assingment_method} is not defined')

    def __finalise__(self):

        """This function is called when all associated fragments have been gathered"""

        # Perfom allele assignment based on likelihood:

        # this is now only generated upon demand, see .allele method

        if self.mapability_reader is not None:

            self.update_mapability()

        self.finalised = True

    def update_mapability(self, set_mq_zero=False):

        """ Update mapability of this molecule.

        mapping qualities are set to 0 if the mapability_reader returns False

        for site_is_mapable

        The mapability_reader can be set when initiating the molecule, or added later.

        Args:

            set_mq_zero(bool) : set mapping quality of associated reads to 0 when the

            mappability reader returns a bad verdict

        Tip:

            Use `createMapabilityIndex.py` to create an index to feed to the mapability_reader

        """

        mapable = None

        try:

            mapable = self.mapability_reader.site_is_mapable(

                *self.get_cut_site())

        except TypeError:

            pass

        except Exception as e:

            raise

        if mapable is False:

            self.set_meta('mp', 'bad')

            if set_mq_zero:

                for read in self.iter_reads():

                    read.mapping_quality = 0

        elif mapable is True:

            self.set_meta('mp', 'unique')

        else:

            self.set_meta('mp', 'unknown')

    def calculate_consensus(self, consensus_model, molecular_identifier, out, **model_kwargs):

        """

        Create consensus read for molecule

        Args:

            consensus_model

            molecular_identifier (str) : identier for this molecule, will be suffixed to the reference_id

            out(pysam.AlingmentFile) : target bam file

            **model_kwargs : arguments passed to the consensus model

        """

        try:

            consensus_reads = self.deduplicate_to_single_CIGAR_spaced(

                out,

                f'c_{self.get_a_reference_id()}_{molecular_identifier}',

                consensus_model,

                NUC_RADIUS=model_kwargs['consensus_k_rad']

            )

            for consensus_read in consensus_reads:

                consensus_read.set_tag('RG', self[0].get_read_group())

                consensus_read.set_tag('mi', molecular_identifier)

                out.write(consensus_read)

        except Exception as e:

            self.set_rejection_reason('CONSENSUS_FAILED', set_qcfail=True)

            self.write_pysam(out)

    def get_a_reference_id(self):

        """

        Obtain a reference id for a random associated mapped read

        """

        for read in self.iter_reads():

            if not read.is_unmapped:

                return read.reference_id

        return -1

    def get_consensus_read(self, target_file,

                           read_name, consensus=None,

                           phred_scores=None,

                           cigarstring=None,

                           mdstring=None,

                           start=None,

                           supplementary=False

                           ):

        """get pysam.AlignedSegment containing aggregated molecule information

        Args:

            target_file(pysam.AlignmentFile) : File to create the read for

            read_name(str) : name of the read to write

        Returns:

            read(pysam.AlignedSegment)

        """

        if start is None:

            start = self.spanStart

        if consensus is None:

            try:  # Obtain the consensus sequence

                consensus = self.get_consensus()

            except Exception as e:

                raise

        if isinstance(consensus, str):

            sequence = consensus

        else:

            sequence = ''.join(

                (consensus.get(

                    (self.chromosome, ref_pos), 'N') for ref_pos in range(

                    self.spanStart, self.spanEnd + 1)))

        if isinstance(phred_scores, dict):

            phred_score_array = list(

                phred_scores.get(

                    (self.chromosome, ref_pos), 0) for ref_pos in range(

                    self.spanStart, self.spanEnd + 1))

        else:

            phred_score_array = phred_scores

        # Construct consensus - read

        cread = pysam.AlignedSegment(header=target_file.header)

        cread.reference_name = self.chromosome

        cread.reference_start = start

        cread.query_name = read_name

        cread.query_sequence = sequence

        cread.query_qualities = phred_score_array

        cread.is_supplementary = supplementary

        if cigarstring is not None:

            cread.cigarstring = cigarstring

        else:

            cread.cigarstring = f'{len(sequence)}M'

        cread.mapping_quality = self.get_max_mapping_qual()

        cread.is_reverse = self.strand

        if mdstring is not None:

            cread.set_tag('MD', mdstring)

        self.write_tags_to_psuedoreads((cread,))

        return cread

    """ method = 1

        sequence = []

        cigar = []

        if method==0:

            prev_end = None

            for block_start,block_end in molecule.get_aligned_blocks():

                if molecule.strand:

                    print(block_end>block_start,block_start, block_end)

                if prev_end is not None:

                    cigar.append(f'{block_start - prev_end}D')

                block_len = block_end-block_start+1

                cigar.append(f'{block_len}M')

                for ref_pos in range(block_start,block_end+1):

                    call = consensus.get((molecule.chromosome, ref_pos),'N')

                    sequence.append(call)

                prev_end = block_end+1

            cigarstring = ''.join(cigar)

        """

    def get_feature_vector(self, window_size=90):

        """ Obtain a feature vector representation of the molecule

        Returns:

            feature_vector(np.array)

        """

        return np.array([

            self.get_strand(),

            self.has_valid_span(),

            self.get_umi_error_rate(),

            self.get_consensus_gc_ratio(),

            len(self.get_raw_barcode_sequences()),

            self.get_safely_aligned_length(),

            self.get_max_mapping_qual(),

            (self.alleles is None),

            self.contains_valid_fragment(),

            self.is_multimapped(),

            self.get_feature_window(window_size=window_size)

        ])

    def get_tag_counter(self):

        """

        Obtain a dictionary with tag -> value -> frequency

        Returns:

            tag_obs (defaultdict(Counter)):

                { tag(str) : { value(int/str): frequency:(int) }

        """

        tags_obs = defaultdict(Counter)

        for tag, value in itertools.chain(

                *[r.tags for r in self.iter_reads()]):

            try:

                tags_obs[tag][value] += 1

            except TypeError:

                # Dont count arrays for example

                pass

        return tags_obs

    def write_tags(self):

        """ Write BAM tags to all reads associated to this molecule

        This function sets the following tags:

            - mI : most common umi

            - DA : allele

            - af : amount of associated fragments

            - rt : rt_reaction_index

            - rd : rt_duplicate_index

            - TR : Total RT reactions

            - ap : phasing information (if allele_resolver is set)

            - TF : total fragments

            - ms : size of the molecule (largest fragment)

        """

        self.is_valid(set_rejection_reasons=True)

        if self.umi is not None:

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

        if self.allele is not None:

            self.set_meta('DA', str(self.allele))

        # Set total amount of associated fragments

        self.set_meta('TF', len(self.fragments) + self.overflow_fragments)

        try:

            self.set_meta('ms',self.estimated_max_length)

        except Exception as e:

            # There is no properly defined aligned length

            pass

        # associatedFragmentCount :

        self.set_meta('af', len(self))

        for rc, frag in enumerate(self):

            frag.set_meta('RC', rc)

            if rc > 0:

                # Set duplicate bit

                for read in frag:

                    if read is not None:

                        read.is_duplicate = True

        # Write RT reaction tags (rt: rt reaction index, rd rt duplicate index)

        # This is only required for fragments which have defined random primers

        rt_reaction_index = None

        for rt_reaction_index, ( (contig, random_primer_start, random_primer_sequence), frags) in enumerate(

                self.get_rt_reactions().items()):

            for rt_duplicate_index, frag in enumerate(frags):

                frag.set_meta('rt', rt_reaction_index)

                frag.set_meta('rd', rt_duplicate_index)

                frag.set_meta('rp', random_primer_start)

        self.set_meta('TR', 0 if (rt_reaction_index is None) else rt_reaction_index + 1)

        if self.allele_resolver is not None:

            self.write_allele_phasing_information_tag()

    def write_tags_to_psuedoreads(self, reads):

        """

        Write molecule information to the supplied reads as BAM tags

        """

        # write methylation tags to new reads if applicable:

        if self.methylation_call_dict is not None:

            self.set_methylation_call_tags(

                self.methylation_call_dict, reads=reads)

        for read in reads:

            read.set_tag('SM', self.sample)

            if hasattr(self, 'get_cut_site'):

                read.set_tag('DS', self.get_cut_site()[1])

            if self.umi is not None:

                read.set_tag('RX', self.umi)

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

                read.set_tag('BC', bc)

                read.set_tag('MI', bc + self.umi)

            # Store total amount of RT reactions:

            read.set_tag('TR', len(self.get_rt_reactions()))

            read.set_tag('TF', len(self.fragments) + self.overflow_fragments)

            if self.allele is not None:

                read.set_tag('DA', self.allele)

        if self.allele_resolver is not None:

            self.write_allele_phasing_information_tag(

                self.allele_resolver, reads=reads)

    def deduplicate_to_single(

            self,

            target_bam,

            read_name,

            classifier,

            reference=None):

        """

        Deduplicate all reads associated to this molecule to a single pseudoread

        Args:

            target_bam (pysam.AlignmentFile) : file to associate the read with

            read_name (str) : name of the pseudoread

            classifier (sklearn classifier) : classifier for consensus prediction

        Returns:

            read (pysam.AlignedSegment) : Pseudo-read containing aggregated information

        """

        # Set all associated reads to duplicate

        for read in self.iter_reads():

            read.is_duplicate = True

        features = self.get_base_calling_feature_matrix(reference=reference)

        # We only use the proba:

        base_calling_probs = classifier.predict_proba(features)

        predicted_sequence = ['ACGT'[i] for i in np.argmax(base_calling_probs, 1)]

        phred_scores = np.rint(-10 * np.log10(np.clip(1 - base_calling_probs.max(1), 0.000000001, 0.999999))).astype(

            'B')

        read = self.get_consensus_read(

            read_name=read_name,

            target_file=target_bam,

            consensus=''.join(predicted_sequence),

            phred_scores=phred_scores)

        read.is_read1 = True

        return read

    def deduplicate_to_single_CIGAR_spaced(

            self,

            target_bam,

            read_name,

            classifier=None,

            max_N_span=300,

            reference=None,

            **feature_matrix_args

    ):

        """

        Deduplicate all associated reads to a single pseudoread, when the span is larger than max_N_span

        the read is split up in multi-segments. Uncovered locations are spaced using N's in the CIGAR.

        Args:

            target_bam (pysam.AlignmentFile) : file to associate the read with

            read_name (str) : name of the pseudoread

            classifier (sklearn classifier) : classifier for consensus prediction

        Returns:

            reads( list [ pysam.AlignedSegment ] )

        """

        # Set all associated reads to duplicate

        for read in self.iter_reads():

            read.is_duplicate = True

        if classifier is not None:

            features, reference_bases, CIGAR, alignment_start, alignment_end = self.get_base_calling_feature_matrix_spaced(

                True, reference=reference, **feature_matrix_args)

            base_calling_probs = classifier.predict_proba(features)

            predicted_sequence = ['ACGT'[i] for i in np.argmax(base_calling_probs, 1)]

            reference_sequence = ''.join(

                [base for chrom, pos, base in reference_bases])

            # predicted_sequence[ features[:, [ x*8 for x in range(4) ] ].sum(1)==0 ] ='N'

            predicted_sequence = ''.join(predicted_sequence)

            phred_scores = np.rint(

                -10 * np.log10(np.clip(1 - base_calling_probs.max(1),

                                       0.000000001,

                                       0.999999)

                               )).astype('B')

        reads = []

        query_index_start = 0

        query_index_end = 0

        reference_position = alignment_start  # pointer to current position

        reference_start = alignment_start  # pointer to alignment start of current read

        supplementary = False

        partial_CIGAR = []

        partial_MD = []

        for operation, amount in CIGAR:

            if operation == 'M':  # Consume query and reference

                query_index_end += amount

                reference_position += amount

                partial_CIGAR.append(f'{amount}{operation}')

            if operation == 'N':

                # Consume reference:

                reference_position += amount

                if amount > max_N_span:  # Split up in supplementary alignment

                    # Eject previous

                    # reference_seq =

                    consensus_read = self.get_consensus_read(

                        read_name=read_name,

                        target_file=target_bam,

                        consensus=predicted_sequence[query_index_start:query_index_end],

                        phred_scores=phred_scores[query_index_start:query_index_end],

                        cigarstring=''.join(partial_CIGAR),

                        mdstring=create_MD_tag(

                            reference_sequence[query_index_start:query_index_end],

                            predicted_sequence[query_index_start:query_index_end]

                        ),

                        start=reference_start,

                        supplementary=supplementary

                    )

                    reads.append(consensus_read)

                    if not supplementary:

                        consensus_read.is_read1 = True

                    supplementary = True

                    # Start new:

                    query_index_start = query_index_end

                    reference_start = reference_position

                    partial_CIGAR = []

                else:

                    partial_CIGAR.append(f'{amount}{operation}')

        reads.append(self.get_consensus_read(

            read_name=read_name,

            target_file=target_bam,

            consensus=''.join(predicted_sequence[query_index_start:query_index_end]),

            phred_scores=phred_scores[query_index_start:query_index_end],

            cigarstring=''.join(partial_CIGAR),

            mdstring=create_MD_tag(

                reference_sequence[query_index_start:query_index_end],

                predicted_sequence[query_index_start:query_index_end]

            ),

            start=reference_start,

            supplementary=supplementary

        ))

        # Write last index tag to last read ..

        if supplementary:

            reads[-1].is_read2 = True

        # Write NH tag (the amount of records with the same query read):

        for read in reads:

            read.set_tag('NH', len(reads))

        return reads

    def extract_stretch_from_dict(self, base_call_dict, alignment_start, alignment_end):

        base_calling_probs = np.array(

            [base_call_dict.get((self.chromosome, pos), ('N', 0))[1] for pos in range(alignment_start, alignment_end)])

        predicted_sequence = [base_call_dict.get((self.chromosome, pos), ('N', 0))[0] for pos in

                              range(alignment_start, alignment_end)]

        predicted_sequence = ''.join(predicted_sequence)

        phred_scores = np.rint(

            -10 * np.log10(np.clip(1 - base_calling_probs,

                                   0.000000001,

                                   0.999999999)

                           )).astype('B')

        return predicted_sequence, phred_scores

    def get_base_confidence_dict(self):

        """

        Get dictionary containing base calls per position and the corresponding confidences

        Returns:

            obs (dict) :  (contig (str), position  (int) ) : base (str) : prob correct (list)

        """

        # Convert (contig, position) -> (base_call) into:

        # (contig, position) -> (base_call, confidence)

        obs = defaultdict(lambda: defaultdict(list))

        for read in self.iter_reads():

            for qpos, rpos in read.get_aligned_pairs(matches_only=True):

                qbase = read.seq[qpos]

                qqual = read.query_qualities[qpos]

                # @ todo reads which span multiple chromosomes

                obs[(self.chromosome, rpos)][qbase].append(1 - np.power(10, -qqual / 10))

        return obs

    def deduplicate_majority(self, target_bam, read_name, max_N_span=None):

        obs = self.get_base_confidence_dict()

        reads = list(self.get_dedup_reads(read_name,

                                     target_bam,

                                     obs={reference_position: phredscores_to_base_call(probs)

                                          for reference_position, probs in obs.items()},

                                     max_N_span=max_N_span))

        self.write_tags_to_psuedoreads([read for read in reads if read is not None])

        return reads

    def generate_partial_reads(self, obs, max_N_span=None):

        CIGAR, alignment_start, alignment_end = self.get_CIGAR()

        query_index_start = 0

        query_index_end = 0

        reference_position = alignment_start  # pointer to current position

        reference_start = alignment_start  # pointer to alignment start of current read

        reference_end = None

        partial_CIGAR = []

        partial_MD = []

        partial_sequence = []

        partial_phred = []

        for operation, amount in CIGAR:

            if operation == 'N':

                if max_N_span is not None and amount > max_N_span:

                    yield reference_start, reference_end, partial_sequence, partial_phred, partial_CIGAR, partial_MD

                    # Clear all

                    partial_CIGAR = []

                    partial_MD = []

                    partial_sequence = []

                    partial_phred = []

                else:

                    # Increment

                    partial_CIGAR.append(f'{amount}{operation}')

                    query_index_start += sum((len(s) for s in partial_sequence))

                reference_position += amount

            elif operation == 'M':  # Consume query and reference

                query_index_end += amount

                if len(partial_CIGAR) == 0:

                    reference_start = reference_position

                start_fetch = reference_position

                reference_position += amount

                reference_end = reference_position

                partial_CIGAR.append(f'{amount}{operation}')

                predicted_sequence, phred_scores = self.extract_stretch_from_dict(obs, start_fetch, reference_end)  # [start .. end)

                partial_sequence.append(predicted_sequence)

                partial_phred.append(phred_scores)

        yield reference_start, reference_end, partial_sequence, partial_phred, partial_CIGAR, partial_MD

    def get_dedup_reads(self, read_name, target_bam, obs, max_N_span=None):

        if self.chromosome is None:

            return None # We cannot perform this action

        for reference_start, reference_end, partial_sequence, partial_phred, partial_CIGAR, partial_MD in self.generate_partial_reads(

                obs, max_N_span=max_N_span):

            consensus_read = self.get_consensus_read(

                read_name=read_name,

                target_file=target_bam,

                consensus=''.join(partial_sequence),

                phred_scores= array('B', np.concatenate(partial_phred)), # Needs to be casted to array

                cigarstring=''.join(partial_CIGAR),

                mdstring=create_MD_tag(

                    self.reference.fetch(self.chromosome, reference_start, reference_end),

                    ''.join(partial_sequence)

                ),

                start=reference_start,

                supplementary=False

            )

            consensus_read.is_reverse = self.strand

            yield consensus_read

    def deduplicate_to_single_CIGAR_spaced_from_dict(

            self,

            target_bam,

            read_name,

            base_call_dict,  # (contig, position) -> (base_call, confidence)

            max_N_span=300,

    ):

        """

        Deduplicate all associated reads to a single pseudoread, when the span is larger than max_N_span

        the read is split up in multi-segments. Uncovered locations are spaced using N's in the CIGAR.

        Args:

            target_bam (pysam.AlignmentFile) : file to associate the read with

            read_name (str) : name of the pseudoread

            classifier (sklearn classifier) : classifier for consensus prediction

        Returns:

            reads( list [ pysam.AlignedSegment ] )

        """

        # Set all associated reads to duplicate

        for read in self.iter_reads():

            read.is_duplicate = True

        CIGAR, alignment_start, alignment_end = self.get_CIGAR()

        reads = []

        query_index_start = 0

        query_index_end = 0

        reference_position = alignment_start  # pointer to current position

        reference_start = alignment_start  # pointer to alignment start of current read

        reference_end = None

        supplementary = False

        partial_CIGAR = []

        partial_MD = []

        partial_sequence = []

        partial_phred = []

        for operation, amount in CIGAR:

            if operation == 'N':

                # Pop the previous read..

                if len(partial_sequence):

                    assert reference_end is not None

                    consensus_read = self.get_consensus_read(

                        read_name=read_name,

                        target_file=target_bam,

                        consensus=''.join(partial_sequence),

                        phred_scores=array('B',np.concatenate(partial_phred)),

                        cigarstring=''.join(partial_CIGAR),

                        mdstring=create_MD_tag(

                            self.reference.fetch(self.chromosome, reference_start, reference_end),

                            ''.join(partial_sequence)

                        ),

                        start=reference_start,

                        supplementary=supplementary

                    )

                    reads.append(consensus_read)

                    if not supplementary:

                        consensus_read.is_read1 = True

                    supplementary = True

                    reference_start = reference_position

                    partial_CIGAR = []

                    partial_phred = []

                    partial_sequence = []

                # Consume reference:

                reference_position += amount

                partial_CIGAR.append(f'{amount}{operation}')

            if operation == 'M':  # Consume query and reference

                query_index_end += amount

                # This should only be reset upon a new read:

                if len(partial_CIGAR) == 0:

                    reference_start = reference_position

                reference_position += amount

                reference_end = reference_position

                partial_CIGAR.append(f'{amount}{operation}')

                predicted_sequence, phred_scores = self.extract_stretch_from_dict(base_call_dict, reference_start,

                                                                                  reference_end)  # [start .. end)

                partial_sequence.append(predicted_sequence)

                partial_phred.append(phred_scores)

        consensus_read = self.get_consensus_read(

            read_name=read_name,

            target_file=target_bam,

            consensus=''.join(partial_sequence),

            phred_scores= array('B',np.concatenate(partial_phred)),

            cigarstring=''.join(partial_CIGAR),

            mdstring=create_MD_tag(

                self.reference.fetch(self.chromosome, reference_start, reference_end),

                ''.join(partial_sequence)

            ),

            start=reference_start,

            supplementary=supplementary

        )

        reads.append(consensus_read)

        if not supplementary:

            consensus_read.is_read1 = True

        supplementary = True

        reference_start = reference_position

        partial_CIGAR = []

        # Write last index tag to last read ..

        if supplementary:

            reads[-1].is_read2 = True

            reads[0].is_read1 = True

        # Write NH tag (the amount of records with the same query read):

        for read in reads:

            read.set_tag('NH', len(reads))

        return reads

    def get_base_calling_feature_matrix(

            self,

            return_ref_info=False,

            start=None,

            end=None,

            reference=None,

            NUC_RADIUS=1,

            USE_RT=True,

            select_read_groups=None):

        """

        Obtain feature matrix for base calling

        Args:

            return_ref_info (bool) : return both X and array with feature information

            start (int) : start of range, genomic position

            end (int) : end of range (inclusive), genomic position

            reference(pysam.FastaFile) : reference to fetch reference bases from, if not supplied the MD tag is used

            NUC_RADIUS(int) : generate kmer features target nucleotide

            USE_RT(bool) : use RT reaction features

            select_read_groups(set) : only use reads from these read groups to generate features

        """

        if start is None:

            start = self.spanStart

        if end is None:

            end = self.spanEnd

        with np.errstate(divide='ignore', invalid='ignore'):

            BASE_COUNT = 5

            RT_INDEX = 7 if USE_RT else None

            STRAND_INDEX = 0

            PHRED_INDEX = 1

            RC_INDEX = 2

            MATE_INDEX = 3

            CYCLE_INDEX = 4

            MQ_INDEX = 5

            FS_INDEX = 6

            COLUMN_OFFSET = 0

            features_per_block = 8 - (not USE_RT)

            origin_start = start

            origin_end = end

            end += NUC_RADIUS

            start -= NUC_RADIUS

            features = np.zeros(

                (end - start + 1, (features_per_block * BASE_COUNT) + COLUMN_OFFSET))

            if return_ref_info:

                ref_bases = {}

            for rt_id, fragments in self.get_rt_reactions().items():

                # we need to keep track what positions where covered by this RT

                # reaction

                RT_reaction_coverage = set()  # (pos, base_call)

                for fragment in fragments:

                    for read in fragment:

                        if select_read_groups is not None:

                            if not read.has_tag('RG'):

                                raise ValueError(

                                    "Not all reads in the BAM file have a read group defined.")

                            if not read.get_tag('RG') in select_read_groups:

                                continue

                        # Skip reads outside range

                        if read is None or read.reference_start > (

                                end + 1) or read.reference_end < start:

                            continue

                        for cycle, q_pos, ref_pos, ref_base in pysamiterators.ReadCycleIterator(

                                read, matches_only=True, with_seq=True, reference=reference):

                            row_index = ref_pos - start

                            if row_index < 0 or row_index >= features.shape[0]:

                                continue

                            query_base = read.seq[q_pos]

                            # Base index block:

                            block_index = 'ACGTN'.index(query_base)

                            # Update rt_reactions

                            if USE_RT:

                                if not (

                                               ref_pos, query_base) in RT_reaction_coverage:

                                    features[row_index][RT_INDEX +

                                                        COLUMN_OFFSET +

                                                        features_per_block *

                                                        block_index] += 1

                                RT_reaction_coverage.add((ref_pos, query_base))

                            # Update total phred score

                            features[row_index][PHRED_INDEX +

                                                COLUMN_OFFSET +

                                                features_per_block *

                                                block_index] += read.query_qualities[q_pos]

                            # Update total reads

                            features[row_index][RC_INDEX + COLUMN_OFFSET +

                                                features_per_block * block_index] += 1

                            # Update mate index

                            features[row_index][MATE_INDEX +

                                                COLUMN_OFFSET +

                                                features_per_block *

                                                block_index] += read.is_read2

                            # Update fragment sizes:

                            features[row_index][FS_INDEX +

                                                COLUMN_OFFSET +

                                                features_per_block *

                                                block_index] += abs(fragment.span[1] -

                                                                    fragment.span[2])

                            # Update cycle

                            features[row_index][CYCLE_INDEX +

                                                COLUMN_OFFSET +

                                                features_per_block *

                                                block_index] += cycle

                            # Update MQ:

                            features[row_index][MQ_INDEX +

                                                COLUMN_OFFSET +

                                                features_per_block *

                                                block_index] += read.mapping_quality

                            # update strand:

                            features[row_index][STRAND_INDEX +

                                                COLUMN_OFFSET +

                                                features_per_block *

                                                block_index] += read.is_reverse

                            if return_ref_info:

                                row_index_in_output = ref_pos - origin_start

                                if row_index_in_output < 0 or row_index_in_output >= origin_end - origin_start + 1:

                                    continue

                                ref_bases[ref_pos] = ref_base.upper()

            # Normalize all and return

            for block_index in range(BASE_COUNT):  # ACGTN

                for index in (

                        PHRED_INDEX,