def cigarTuplesToOperations(tuples, includeHardClip=True):
    """
    Produce a sequence of CIGAR operations given a cigar tuples list
    from a pysam AlignedRead instance.

    See https://pysam.readthedocs.io/en/latest/
        api.html#pysam.AlignedSegment.cigartuples

    @param tuples: A C{list} of 2-tuples, each containing an operation
        and an C{int} count.
    @param includeHardClip: If C{True}, yield hard clipping operations,
        else elide them.
    """
    for operation, count in tuples:
        if operation != CHARD_CLIP or includeHardClip:
            for _ in range(count):
                yield operation def softClippedOffset(offset, pairs, cigarOperations):
    """
    Get the offset in the reference of a base that has been soft-clipped
    in a query.

    @param offset: The offset (in C{pairs}) of the soft-clipped base (which
       has a to-be-determined offset in the reference).
    @param pairs: A C{list} of C{int} (queryOffset, referenceOffset) pairs
       for a pysam C{pysam.AlignedSegment}, as returned by
       C{read.get_aligned_pairs} where the read is returned by the pysam
       C{fetch} method.
    @param cigarOperations: A C{list} of pysam CIGAR operations.
    @return: The C{int} reference offset of the soft-clipped base.
    """
    assert cigarOperations[offset] == CSOFT_CLIP
    assert pairs[offset][1] is None

    # Look back.
    count = 0
    for pair, cigarOperation in zip(pairs[offset::-1],
                                    cigarOperations[offset::-1]):
        if cigarOperation == CSOFT_CLIP:
            count += 1
        else:
            _, referenceOffset = pair
            # If we have a reference offset, then we've found our
            # answer. Otherwise, this has to be a hard-clipped base.
            if referenceOffset is None:
                assert cigarOperation == CHARD_CLIP
                # We could break but won't, as explained above.
            else:
                return referenceOffset + count

    # Look ahead.
    count = 0
    for pair, cigarOperation in zip(pairs[offset:], cigarOperations[offset:]):
        if cigarOperation == CSOFT_CLIP:
            count += 1
        else:
            _, referenceOffset = pair
            # If we have a reference offset, then we've found our
            # answer. Otherwise, this has to be a hard-clipped base.
            if referenceOffset is None:
                assert cigarOperation == CHARD_CLIP
                # We could break here, but if we keep going we can test
                # that the whole rest of the CIGAR is hard clips and will
                # fail if not. That's not needed in a world of perfect code.
            else:
                return referenceOffset - count

    # This should be impossible.
    raise ValueError('Soft-clipped base with no following or preceding '
                     'non-hard-clipped bases.') def insertionOffset(offset, pairs, cigarOperations):
    """
    Get the insertion offset in the reference for a base that is part of an
    insertion.

    @param offset: The offset (in C{pairs}) of the soft-clipped base (which
       has a to-be-determined offset in the reference).
    @param pairs: A C{list} of C{int} (queryOffset, referenceOffset) pairs
       for a pysam C{pysam.AlignedSegment}, as returned by
       C{read.get_aligned_pairs} where the read is returned by the pysam
       C{fetch} method.
    @param cigarOperations: A C{list} of pysam CIGAR operations.
    @return: A 2-tuple containing a C{bool} and the C{int} reference offset of
        the base after the insertion. The C{bool} indicates whether we needed
        to look back in the sequence to find a reference base.
    """
    assert cigarOperations[offset] == CINS
    assert pairs[offset][1] is None

    # Look back.
    for pair, cigarOperation in zip(pairs[offset::-1],
                                    cigarOperations[offset::-1]):
        if cigarOperation != CINS:
            _, referenceOffset = pair
            # If we have a reference offset, then we've found our
            # answer. Otherwise, check as above.
            if referenceOffset is None:
                assert cigarOperation not in CONSUMES_REFERENCE
            else:
                return True, referenceOffset + 1

    # Look ahead.
    for pair, cigarOperation in zip(pairs[offset:], cigarOperations[offset:]):
        if cigarOperation != CINS:
            _, referenceOffset = pair
            # If we have a reference offset, then we've found our
            # answer. Otherwise, check that this is something that does not
            # consume reference bases.
            if referenceOffset is None:
                assert cigarOperation not in CONSUMES_REFERENCE
            else:
                return False, referenceOffset

    # This should be impossible.
    raise ValueError('Inserted base with no following or preceding '
                     'reference bases.') class Bases:
    """
    Manage a collection of (base, quality) pairs for a genome site.
    """
    __slots__ = ('count', 'counts')

    def __init__(self):
        self.count = 0
        self.counts = dict.fromkeys('ACGT', 0)

    def __eq__(self, other):
        return self.count == other.count and self.counts == other.counts

    def __str__(self):
        return f'<Bases count={self.count}, bases={self.counts}'

    __repr__ = __str__

    def __getitem__(self, key):
        return self.counts[key]

    def __len__(self):
        return self.count

    def __add__(self, other):
        new = Bases()
        new.count = self.count + other.count
        for counts in self.counts, other.counts:
            for key, count in counts.items():
                new.counts[key] += count
        return new

    def __iadd__(self, other):
        self.count += other.count
        for key, count in other.counts.items():
            self.counts[key] += count
        return self

    def append(self, base, quality):
        """
        Append a (base, quality) pair.

        @param base: A C{str} nucleotide base.
        @param quality: An C{int} nucleotide quality score.
        @return: C{self}.
        """
        if base != 'N':
            self.count += 1
            self.counts[base] += quality

        return self

    def consensus(self, threshold, minCoverage, lowCoverage, noCoverage):
        """
        Get the base that can be used as part of a consensus.

        If there are sufficient reads, this is the least-ambiguous nucleotide
        code for our bases, given a required homogeneity threshold. Otherwise,
        the low coverage value.

        @param threshold: A C{float} threshold. This fraction, at least, of the
            most-common nucleotides at a site are used to determine the
            consensus nucleotide (or ambiguous symbol if more than one
            nucleotide is required to achieve this threshold). If there is a
            tie in nucleotide counts at a site that causes the threshold to be
            met, all nucleotides of equeal frequncy will be included in the
            ambiguous symbol for that site. This is perhaps better explained
            with an example. See
            https://assets.geneious.com/manual/2020.1/static/GeneiousManualse43.html
            and the corresponding testGeneiousExamplesTie test in
            test/test_dna.py
        @param minCoverage: An C{int} minimum number of reads that must cover a
            site for a consensus base to be called. If fewer reads cover a
            site, the C{lowCoverage} value is used.
        @param lowCoverage: A C{str} indicating what base to use when
            0 < N < minCoverage reads cover a site.
        @param noCoverage: A C{str} indicating what base to use when
            no reads cover the site.
        @return: A C{str} nucleotide code. This will be an ambiguous code if
            the homogeneity C{threshold} is not met.
        """
        return (noCoverage if self.count == 0 else
                (lowCoverage if self.count < minCoverage else
                 leastAmbiguousFromCounts(self.counts, threshold))) @contextmanager
def maybeProgressBar(show, maxValue, prefix):
    """
    A context manager to maybe show a progress bar.

    @param show: If C{True}, yield a progress bar, else a Class with an
        C{update} method that does nothing.
    @param maxValue: The C{int} number of tasks to show progress for.
    @param prefix: A C{str} prefix, to appear at the start of the progress bar.
    """
    if show and os.isatty(2):
        widgets = [
            progressbar.SimpleProgress(format='%(value_s)s/%(max_value_s)s'),
            progressbar.Percentage(format=' %(percentage)3d%%'),
            ' ',
            progressbar.Bar(marker='\x1b[33m#\x1b[39m'),
            ' ',
            progressbar.Timer(format='Elapsed: %(elapsed)s'),
            ' ',
            progressbar.ETA(format='ETA: %(eta)8s'),
        ]
        with progressbar.ProgressBar(max_value=maxValue, widgets=widgets,
                                     prefix=prefix) as bar:
            yield bar
    else:
        class Bar:
            update = staticmethod(lambda _: None)
        yield Bar class DNARead(_NucleotideRead):
    """
    Hold information and methods to work with DNA reads.
    """
    ALPHABET = set('ATCG')

    COMPLEMENT_TABLE = _makeComplementTable(ambiguous_dna_complement) @contextmanager
def samfile(filename):
    """
    A context manager to open and close a SAM/BAM file.

    @param filename: A C{str} file name to open.
    """
    f = AlignmentFile(filename)
    yield f
    f.close() def getReferenceInfo(bam, bamFilename, bamId=None, referenceFasta=None,
                     fastaId=None, quiet=False):
    """
    Get the id and length of a BAM file reference seqeunce.

    @param bam: An open BAM file.
    @param bamFilename: The C{str} file name of the BAM file.
    @param bamId: A C{str} BAM file reference name indicating which aligned
        reads to make a consensus from. If not given, will be inferred
        from the BAM file header.
    @param referenceFasta: A C{str} file name containing the sequence that was
        aligned to in making the BAM file.
    @param fastaId: A C{str} reference name indicating which sequence in
        C{referenceFasta} to use as a reference. Only considered if
        C{referenceFasta} is given. If not given and C{referenceFasta} is,
        the reference id will be inferred from reference names in the BAM
        header, or will be taken as the id of the first sequence in
        C{referenceFasta}.
    @param quiet: If C{True}, suppress diagnostic output.
    @raise UnknownReference: For an unknown reference.
    @raise UnspecifiedReference: If a reference is not given and cannot be
        inferred.
    @raise ReferenceNameMismatchError: If the names of the reference in the BAM
        file and the FASTA reference file (if given) do not match (unless both
        are explicitly given, in which case mismatching names are allowed as
        long as the sequence lengths are the same).
    @raise UnequalReferenceLengthError: If the lengths of the reference in the
        BAM file and the FASTA reference file (if given) do not match
    @return: A 2-tuple of C{str}, C{int} giving the refrence sequence id and
        length.
    """
    bamReferences = set(samReferences(bam))
    checkName = checkLength = True
    inferredBamId = bamId is None
    inferredFastaId = fastaId is None
    referenceLength = None

    if bamId is None:
        inferredBamId = True
    else:
        if bamId not in bamReferences:
            raise UnknownReference(
                f'BAM file {str(bamFilename)!r} does not mention a reference '
                f'with id {bamId!r}. Known references are: '
                f'{", ".join(sorted(bamReferences))}.')

        inferredBamId = False
        tid = bam.get_tid(bamId)
        referenceLength = bam.lengths[tid]

        if fastaId is not None:
            # Both ids have been given explicitly, so don't check that they
            # are identical, assume the user knows what they are doing (we
            # will just check the lengths, if necessary).
            checkName = False

    if referenceFasta is None:
        reference = None
    else:
        if fastaId is None:
            # We're given a reference FASTA file, but no id to use. Look at
            # the sequences in the FASTA and take the first one that has a
            # name and length matching a BAM reference (we could actually
            # look to see if there are more than one and, if so, give an
            # error if their sequences are not identical).
            firstRead = None
            for read in FastaReads(referenceFasta):
                if firstRead is None:
                    firstRead = read
                if read.id in bamReferences:
                    tid = bam.get_tid(read.id)
                    if len(read) == bam.lengths[tid]:
                        reference = read
                        checkLength = checkName = False
                        break
            else:
                # Try the first sequence in the FASTA file.
                reference = firstRead

            if firstRead is None:
                raise UnknownReference(
                    f'The FASTA reference file {str(referenceFasta)!r} '
                    f'contained no sequences.')
        else:
            for read in FastaReads(referenceFasta):
                if read.id == fastaId:
                    reference = read
                    break
            else:
                raise UnknownReference(f'No sequence with id {fastaId!r} '
                                       f'found in {str(referenceFasta)!r}.')

        # Set reference length if it has not already been set due to a
        # given bamId.
        if referenceLength is None:
            referenceLength = len(reference)

    if bamId is None:
        if len(bamReferences) == 1:
            # This is the only possibility.
            bamId = tuple(bamReferences)[0]
            tid = bam.get_tid(bamId)
            referenceLength = bam.lengths[tid]
        elif reference and reference.id in bamReferences:
            tid = bam.get_tid(reference.id)
            if len(reference) == bam.lengths[tid]:
                bamId = reference.id
                checkLength = checkName = False
        else:
            raise UnspecifiedReference(
                f'Could not infer a BAM reference. Available references are: '
                f'{", ".join(sorted(bamReferences))}.')

    if inferredBamId and not quiet:
        print(f'BAM reference id {bamId!r} inferred from context.',
              file=sys.stderr)

    if reference:
        if inferredFastaId and not quiet:
            print(f'FASTA reference id {reference.id!r} inferred from '
                  f'context.', file=sys.stderr)

        if checkName and reference.id != bamId:
            raise ReferenceNameMismatchError(
                f'Reference FASTA sequence name {reference.id!r} does '
                f'not match the BAM id {bamId!r}.')

        if checkLength and len(reference) != referenceLength:
            raise UnequalReferenceLengthError(
                f'Reference FASTA sequence {reference.id!r} has length '
                f'{len(reference)}, but the BAM reference {bamId!r} has '
                f'length {referenceLength}.')

    return bamId, reference, referenceLength class UnspecifiedReference(SamError):
    "Reference sequence not specified." CONSUMES_REFERENCE = {CMATCH, CDEL, CREF_SKIP, CEQUAL, CDIFF} def pct(a, b):
    """
    Format a string showing two integers and what percentage the first
    is of the second.

    @param a: An C{int}, the numerator.
    @param b: An C{int}, the denominator.
    """
    assert 0 <= a <= b
    if b:
        return ('%d/%d (%.3f%%)' %
                (a, b, (a / b if b else 0.0) * 100.0))
    else:
        return '0/0 (0.000%)' @contextmanager
def openOr(filename, mode='r', defaultFp=None, specialCaseHyphen=True):
    """
    A context manager to either open a file or use a pre-opened default file.

    @param filename: If not C{None}, this is the argument to pass to C{open}
        along with C{mode}. If C{None}, C{fp} is used.
    @param mode: The C{str} file opening mode, used if C{filename} is not
        C{None}.
    @param defaultFp: An open file-like object to yield if C{filename} is
        C{None}.
    @param specialCaseHyphen: If C{True}, treat '-' as a C{None} filename
        and yield the C{defaultFp}.
    """
    if filename is None or filename == '-' and specialCaseHyphen:
        yield defaultFp
    else:
        with open(filename, mode) as fp:
            yield fp
import sys
from collections import defaultdict
from pysam import CINS, CDEL, CSOFT_CLIP
from dark.cigar import (
    cigarTuplesToOperations, softClippedOffset, insertionOffset)
from dark.dna import Bases
from dark.progress import maybeProgressBar
from dark.reads import DNARead
from dark.sam import (
    samfile, getReferenceInfo, UnspecifiedReference, CONSUMES_REFERENCE)
from dark.utils import pct, openOr


def getPairs(bam, bamId, referenceLength, ignoreQuality,
             includeSoftClipped, progress):
    """Compute a majority consensus using pysam's fetch function.

    @param bam: An open BAM file.
    @param bamId: A C{str} reference name indicating which reference to
        find reads against in the BAM file. If C{None} and there is only one
        reference in the BAM file, that one will be used, else a ConsensusError
        is raised due to not knowing which reference to use.
    @param reference: A C{Read} instance giving the reference sequence, or
        C{None}.
    @param referenceLength: The C{int} length of the reference (note that we
        may not have the reference sequence but we can still get its length
        from the BAM file header).
    @param ignoreQuality: If C{True}, ignore quality scores.
    @param includeSoftClipped: If C{True} include information from read bases
        marked as soft-clipped by the algorithm that made the BAM file.
    @param progress: If C{True}, display a progress bar on standard error.
    @return: A 3-tuple of correspondences, deletions, insertions, with types
        as in the lines just below.
    """
    correspondences = defaultdict(lambda: Bases())
    deletions = defaultdict(int)
    insertions = {}
    nReads = bam.count(contig=bamId)

    with maybeProgressBar(progress, nReads, 'Reads    : ') as bar:
        for readCount, read in enumerate(bam.fetch(contig=bamId), start=1):
            if not read.is_unmapped:
                addPairsInfo(
                    read.get_aligned_pairs(),
                    list(cigarTuplesToOperations(read.cigartuples,
                                                 includeHardClip=False)),
                    read.query_name,
                    read.query_sequence,
                    ([1] * len(read.query_sequence) if ignoreQuality else
                     read.query_qualities), referenceLength,
                    includeSoftClipped, correspondences, deletions, insertions)

            bar.update(readCount)

    return correspondences, deletions, insertions


def addPairsInfo(pairs, cigarOperations, queryId, query, qualities,
                 referenceLength, includeSoftClipped, correspondences,
                 deletions, insertions):
    """
    Add information about matched pairs of nucleotides.

    @param pairs: A C{list} of 2-C{tuple}s of query offset, reference offset.
        Either (but not both) member of each tuple might be C{None} to indicate
        an indel mismatch.
    @param cigarOperations: A C{list} of CIGAR operations corresponding to the
        information in C{pairs}.
    @param queryId: A C{str} query id.
    @param query: A C{str} query DNA sequence.
    @param qualities: A C{list} of quality scores.
    @param includeSoftClipped: Include information from read bases that were
        marked as soft-clipped by the algorithm that made the BAM file.
    @param correspondences: A C{defaultdict(list)}, to hold (base, quality)
        scores for when a query offset corresponds to a reference offset.
    @param deletions: A C{set} of C{int} reference offsets that are deleted in
        the query.
    @param insertions: A C{defaultdict(list)}, to hold (base, quality)
        scores for when a query contains an insertion to the reference.
    """
    if len(pairs) != len(cigarOperations):
        raise ValueError(
            f'Query {queryId!r} (length {len(query)}) with {len(pairs)} pairs '
            f'but {len(cigarOperations)} CIGAR operaions:'
            f'\nPairs: {pairs}\nCIGAR: {cigarOperations}')

    assert not any(pair == (None, None) for pair in pairs)

    inInsertion = False

    for count, ((queryOffset, referenceOffset), cigarOperation) in enumerate(
            zip(pairs, cigarOperations)):

        if queryOffset is None:
            # The query is missing something that is in the reference. So this
            # is a deletion from the reference.
            assert cigarOperation == CDEL
            assert referenceOffset is not None
            deletions[referenceOffset] += 1
            inInsertion = False

        elif referenceOffset is None:
            base = query[queryOffset]
            quality = qualities[queryOffset]

            if cigarOperation == CINS:
                # The query has an insertion (relative to the reference).

                # A CIGAR string shouldn't start with an insertion, IMO.
                # Rather, in such a case, it must start with unmatched
                # (soft-clipped) bases.
                # assert lastReferenceOffset is not None

                lookedBack, iOffset = insertionOffset(count, pairs,
                                                      cigarOperations)
                if not inInsertion:
                    inInsertion = True
                    if iOffset not in insertions:
                        insertions[iOffset] = Insertion(iOffset)
                    insertions[iOffset].start(iOffset if lookedBack else None)

                insertions[iOffset].append(base, quality)
            else:
                assert cigarOperation == CSOFT_CLIP
                inInsertion = False
                if includeSoftClipped:
                    correspondences[
                        softClippedOffset(count, pairs, cigarOperations)
                    ].append(base, quality)
        else:
            # Query and reference offsets are both non-None.

            assert cigarOperation in CONSUMES_REFERENCE