from singlecellmultiomics.molecule.molecule import Molecule
from singlecellmultiomics.utils.sequtils import phred_to_prob
import numpy as np

class ScarTraceMolecule(Molecule):
    """ScarTrace Molecule class

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

        **kwargs: extra args

    """

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

        # Run code to obtain scar description
        assert len(self) == 1, 'only implemented for molecules with a single associated fragment'

        scarDescription = set()
        qualities = []
        for read in self.iter_reads():
            if read is None:
                continue
            if read.is_unmapped:
                continue
            firstCigarOperation, firstCigarLen = read.cigartuples[0]
            insertPos = 0
            lastNonNoneRefPos = read.reference_start if firstCigarOperation != 4 else read.reference_start - firstCigarLen

            expandedCigar = []
            for cigarOperation, cigarLen in read.cigartuples:
                expandedCigar += [cigarOperation] * cigarLen

            for queryPos, referencePos in read.get_aligned_pairs(
                    matches_only=False):
                if queryPos is not None:
                    qualities.append(phred_to_prob(read.qual[queryPos]))

                if queryPos is None and referencePos is None:
                    continue

                if referencePos is not None:
                    lastNonNoneRefPos = referencePos
                    insertPos = 0
                    if queryPos is not None:  # If both the reference and query match, there is not scar information
                        continue

                if queryPos is None:
                    scarDescription.add(f'{referencePos}.D')
                elif referencePos is None:  # insert or clip:
                    operation = expandedCigar[queryPos]
                    if operation == 1:
                        if lastNonNoneRefPos is None:
                            raise ValueError('Unsolvable :(')
                        queryBase = read.seq[queryPos]
                        scarDescription.add(
                            f'{queryBase}.{insertPos+lastNonNoneRefPos}.I')
                        insertPos += 1

        scarDescription = ','.join(sorted(list(scarDescription)))

        # Add average base calling quality excluding primers:
        meanQ = np.mean(qualities) if len(qualities)>0 else 0
        for read in self.iter_reads():
            read.set_tag('SQ', 1 - meanQ)
            if len(scarDescription) == 0:
                scarDescription = 'WT'
            #@todo when collapsing into molecule:
            read.set_tag('SD', scarDescription)
            #@todo when collapsing into molecule:
        if self[0].get_R1() is not None:
            self.set_meta('DS', self[0].get_R1().reference_start)
            self.set_meta('RZ', 'SCAR')
            self.set_meta('DT', 'SCAR')


    def write_tags(self):
        Molecule.write_tags(self)

    def is_valid(self, set_rejection_reasons=False, reference=None):
        if not any( fragment.is_valid() for fragment in self ):
            self.set_rejection_reason('all_fragments_rejected')
            return False

        return True