from collections import Counter

from dark.errors import NoSuchGenomeError
from dark.genbank import GenomeRanges
from dark.reads import DNARead
from dark.sam import samfile


class GenomeProteinInfo(object):
    """
    Hold information about the proteins in a genome and how they are matched
    by reads in SAM files.

    @param accession: The C{str} accession number of a genome.
    @param proteinGenomeDB: A L{dark.civ.proteins.SqliteIndex} instance.
    @param checkTranslations: If C{True}, check that the protein sequences
        that area supposed to have come from the genome can be obtained by
        translating the corresponding region of the genome.
    """
    def __init__(self, genomeAccession, proteinGenomeDB,
                 checkTranslations=True):
        self.genomeAccession = genomeAccession
        self.proteinGenomeDB = proteinGenomeDB
        # self.proteins is keyed by protein accession number.
        self.proteins = {}
        self.coveredProteins = set()
        # self.offsets is keyed by offset, values are dicts that contain a list
        # of protein accession numbers that overlap that offset and a set of
        # read ids (if any) that match at that offset. The offsets keys are
        # only those that correspond to a protein.
        self.offsets = {}
        # self.coveredOffsetCount holds the read counts for all offsets covered
        # by reads, regardless of whether the offsets correspond to proteins or
        # not.
        self.coveredOffsetCount = Counter()
        self.samFiles = []
        self.readIdsMatchingGenome = set()

        self.genome = proteinGenomeDB.findGenome(genomeAccession)
        if self.genome is None:
            raise NoSuchGenomeError('Reference %r not found in protein/genome '
                                    'database.' % genomeAccession)

        for protein in proteinGenomeDB.findProteinsForGenome(genomeAccession):
            proteinAccession = protein['accession']
            self.proteins[proteinAccession] = protein

            ranges = GenomeRanges(protein['offsets']).ranges
            # print('Protein accession', proteinAccession)
            # print(ranges)

            for (start, stop, forward) in ranges:
                for offset in range(start, stop):
                    if offset not in self.offsets:
                        self.offsets[offset] = {
                            'proteinAccessions': set(),
                            'readIds': set(),
                        }
                    self.offsets[offset]['proteinAccessions'].add(
                        proteinAccession)

            if checkTranslations:
                self._checkTranslation(self.genome, ranges, protein)

    def _checkTranslation(self, genome, ranges, protein):
        """
        Make sure all protein sequences supposed to be in the genome can in
        fact be obtained by translating the genome.

        @param genome: A C{dict} with genome information from our sqlite3
            protein/genome database, as returned by
            C{dark.civ.proteins.SqliteIndex.findGenome.
        @param ranges: A C{list} of (start, stop, forward) nucleotide ranges
            for the protein in the genome.
        @param protein: A C{dict} with protein information from our sqlite3
            protein/genome database, as returned by
            C{dark.civ.proteins.SqliteIndex.findProtein.
        """
        proteinSequence = protein['sequence'] + '*'

        # print('protein name', protein['product'], 'ranges', ranges)
        sequence = ''.join([genome['sequence'][start:stop]
                            for (start, stop, _) in ranges])

        genomeRead = DNARead('id', sequence)
        translations = list(genomeRead.translations())
        index = 0 if protein['forward'] else 3
        if translations[index].sequence != proteinSequence:
            # TODO: improve this error to show what actually went wrong.
            raise ValueError(
                'Could not translate genome range to get protein sequence')

    def addSAM(self, filename, filterAlignment=None):
        """
        Read a SAM file and add information about the reads that match our
        reference id.

        @param filename: A C{str} SAM filename.
        @param filterAlignment: A 1-argument function to be used for filtering
            reads in the SAM file. If C{None}, all alignments will be examined.
        """
        self.samFiles.append(filename)
        referenceId = self.genomeAccession
        with samfile(filename) as sam:
            for column in sam.pileup():
                for read in column.pileups:
                    alignment = read.alignment
                    if (alignment.reference_name == referenceId and
                            (filterAlignment is None or
                             filterAlignment(alignment))):
                        readId = alignment.query_name
                        self.readIdsMatchingGenome.add(readId)
                        offset = column.reference_pos
                        self.coveredOffsetCount[offset] += 1

                        try:
                            offsetInfo = self.offsets[offset]
                        except KeyError:
                            pass
                        else:
                            # This offset corresponds to one or more proteins.
                            self.coveredProteins.update(
                                offsetInfo['proteinAccessions'])
                            offsetInfo['readIds'].add(readId)

    def proteinCoverageInfo(self, proteinAccession):
        """
        Calculate coverage information for a protein.

        @param proteinAccession: A C{str} accession number.
        @raises KeyError: If C{proteinAccession} is not known.
        @return: A C{dict} containing
                * the number of covered offsets,
                * the total number of read bases that cover the protein,
                * the protein length (in nucleotides), and
                * the set of all matching read ids.
            See below for the dictionary keys.
        """
        protein = self.proteins[proteinAccession]
        coveredOffsets = 0
        totalBases = 0
        allReadIds = set()
        offsetsSeen = set()
        proteinLength = 0

        for (start, stop, forward) in GenomeRanges(protein['offsets']).ranges:
            proteinLength += stop - start
            for offset in range(start, stop):
                assert offset not in offsetsSeen
                offsetsSeen.add(offset)
                readIds = self.offsets[offset]['readIds']
                if readIds:
                    coveredOffsets += 1
                    totalBases += len(readIds)
                    allReadIds.update(readIds)

        # The +3 in the following is because the database holds the AA length,
        # not including the stop codon. But the database range covers the stop
        # codon.
        dbProteinLength = self.proteins[proteinAccession]['length'] * 3 + 3
        if proteinLength != dbProteinLength:
            raise ValueError(
                'Sum of protein database ranges (%d) does not agree with '
                'database protein length (%d) for protein %s!' %
                (proteinLength, dbProteinLength, proteinAccession))

        return {
            'coveredOffsets': coveredOffsets,
            'totalBases': totalBases,
            'ntLength': proteinLength,
            'readIds': allReadIds,
        }

    def readIdsForAllProteins(self):
        """
        Get the set of read ids for reads matching any protein at all.

        @return: A C{set} of C{str} read ids.
        """
        result = set()
        for offsetInfo in self.offsets.values():
            result.update(offsetInfo['readIds'])
        return result
