#include "Genome_Index_TableQ.h"

CGenome_Index_TableQ::CGenome_Index_TableQ(void)
{
    initialization();
}

CGenome_Index_TableQ::~CGenome_Index_TableQ(void)
{
    // this->pgenomeNTInBits and this->pgenome will be deleted in parent class;
}

int CGenome_Index_TableQ::initialization(void)
{
    this->bExcludeAmbiguous = true;
    return(0);
}

bool CGenome_Index_TableQ::getSeqFromFasta(const char* genomeListfileName, string refFormat)
{
    getBasename(genomeListfileName, this->caRefName);
    if (fileExist(genomeListfileName)) {
        this->getGenomeNTdata(genomeListfileName, refFormat);
        TIME_INFO(this->pgenomeNTInBits = new CGenomeInBits(this->pgenomeNT), "Build genome in bits ");
        this->pgenomeNT->freeChromosomeSpace();
        return(true);
    } else {
        cout << "File " << genomeListfileName << " not found!" << endl;
        return(false);
    }
}

bool CGenome_Index_TableQ::getSeqFromDS(CGenomeNTdata* pgenomeNT)
{
    this->pgenomeNT = pgenomeNT;
    myStrCpy(this->caRefName, pgenomeNT->refName, FILENAME_MAX);
    TIME_INFO(this->pgenomeNTInBits = new CGenomeInBits(this->pgenomeNT), "Build genome in bits ");
    this->pgenomeNT->freeChromosomeSpace();
    return(true);
}

// (1) Query each seed pattern for hit. (2) Check each hit for valid alignment and put in the palignmentsQ
// (3) Return the minDiff for alignment
pair<CIndex_Type*, CIndex_Type*> CGenome_Index_TableQ::queryKmer\
(CReadInBits window, unsigned int shift) const
{
    unsigned int BIN_SIZE_CHECK_THRESHOLD = 1; // If there are only this number in bin, check

    window = window.getSuffixStr(shift);

    unsigned int uiHashValue = this->getHashValue(window);
    unsigned int uiSeedKeyL = this->getSeedKey(window);
    unsigned int uiSeedKeyU;
    if (bEXTEND_SEED) {
        uiSeedKeyU = this->getSeedKeyUpperBound(window, shift);
    } else {
        uiSeedKeyU = uiSeedKeyL;
    }
    // cout << uiHashValue << ',' << uiSeedKey << endl;

    unsigned int bucketStart = this->pHashIndexTable->aiIndexTable[uiHashValue];
    unsigned int nextBucketStart = this->pHashIndexTable->aiIndexTable[uiHashValue + 1];

    if (bucketStart > nextBucketStart) ERR; // Error Case
    if (bucketStart >= nextBucketStart) { // empty bucket
        CIndex_Type* tmp = NULL;
        return(pair<CIndex_Type*, CIndex_Type*>(tmp, tmp)); //No hit
    }

    CIndex_Type* hitStart = &this->pIndexTable[bucketStart];
    CIndex_Type* hitEndPlus1 = &this->pIndexTable[nextBucketStart];

    // If the bin is large enough, use binary search to find the correct hituiSeedKey
    if (bucketStart + BIN_SIZE_CHECK_THRESHOLD  < nextBucketStart) {
        hitStart    = lower_bound(hitStart, hitEndPlus1, uiSeedKeyL, CcompareFunctor4LowerBound(this));
        // TODO fill the suffix digit for the uiSeedKey
        hitEndPlus1 = upper_bound(hitStart, hitEndPlus1, uiSeedKeyU, CcompareFunctor4UpperBound(this));
    } // else assume everything in the bucket is a hit without equal_range check
    return(pair<CIndex_Type*, CIndex_Type*>(hitStart, hitEndPlus1));
}

unsigned int CGenome_Index_TableQ::getSeedKeyUpperBound(CReadInBits window, unsigned int shift) const
{
    //WORD_SIZE padding = 0xffffffffffffffff;
    WORD_SIZE padding = -1;
    if (this->bMapReadInColors) {
        padding <<= (this->uiRead_Length - 1 - shift);

    } else {
        padding <<= (this->uiRead_Length - shift);
    }
    window.UpperBits |= padding;
    window.LowerBits |= padding;

    return(this->getSeedKey(window));
}

// Query a read in bases (Illumina) for hit and check uiDiff and put the result into the given Queue
unsigned int CGenome_Index_TableQ::queryReadBases(CReadInBits readInBases, CAlignmentsQ& aQue, bool bClearQ, bool bForward) const
{
    if (bClearQ) {
        aQue.clearHits();
    }
    if (!bForward) { //Query reverse complement read for alignment
        readInBases = reverseCompliment(this->uiRead_Length, readInBases);
    }
    bool bUseShortCut = this->bExcludeAmbiguous && !(aQue.qAllInThreshold());
    bool qAllHits = aQue.qAllInThreshold();
    for (unsigned int shift = 0; shift <= this->uiNoOfShift; shift++) {
        // If the best match is exact matched, no need to go to the next shift
        // for all exact matches must be found previously.
        if ((aQue.MinDiff == 0) && (shift > 0) && (aQue.load > 0) && !qAllHits) {
            break;
        }
        pair<CIndex_Type*, CIndex_Type*> hits = queryKmer(readInBases, shift);
        if (hits.first != NULL) {
            for (CIndex_Type* it = hits.first; it < hits.second; it++) {
                if (*it >= shift) {
                    unsigned int alignStart = *it - shift;
                    if (isMasked(alignStart)) {
                        continue;
                    } else {
                        CReadInBits ref = this->pgenomeNTInBits->getSubstringInBits(alignStart);
                        unsigned int uiDiff = bitsStrNCompare(ref, readInBases, this->uiRead_Length);
                        // The flag in alignmentsQ decide whether all alignments within uiDiff or
                        // only the alignments with minimum Diff are queue.
                        if (uiDiff <= this->uiSubDiffThreshold) {
                            aQue.saveHits(alignStart, (unsigned short)uiDiff);
                        }

                        if (bUseShortCut) { // short cut to exclude ambiguous reads
                            bool bMap2Repeat = this->pbaRepeatRepresentativeFlag->b(alignStart);
                            if (bMap2Repeat) {
                                aQue.AmbiguousFlag = true;
                            }
                            if (aQue.MinDiff == 0 && (aQue.load >= 2 || bMap2Repeat)) {
                                aQue.setForwardLoad(bForward);
                                return(0);
                            }
                        }
                    }
                }
            }
            if (aQue.MinDiff == 0) {
                // short cut. Output no more than iMaxCapacity exaxt alignment
                /* if (aQue.load >= aQue.iMaxCapacity - 1) {
                    aQue.setForwardLoad(bForward);
                    return(0);
                }*/
                // short cut. All exact matches will be found after first shift
                if (bUseShortCut && aQue.load >= 2) {
                    aQue.setForwardLoad(bForward);
                    return(1);
                }
            }
        }
    }
    /* DEBUG
    if(palignmentsQ.MinDiff <= this->uiSubDiffThreshold) {
        CReadInBits ref = this->pgenomeNTInBits->getSubstringInBits(aQue.aiHitIndex[0]);
        printBitsStrCompare(ref, readInBases, "Good Alignments!!");
    }*/

    // The records before are for the forward direction if this is a forward query
    aQue.setForwardLoad(bForward);
    return(aQue.MinDiff);
}

// Query a read in bases (Illumina) for hit and check uiDiff and put the result into the given Queue
unsigned int CGenome_Index_TableQ::queryLongReadBases(CReadInBits r1, CReadInBits r2, bool oddReadLength, CAlignmentsQ& aQue, int queryHalf, bool bClearQ, bool bForward) const
{
    const unsigned int firstPartLength = this->uiRead_Length;
    const unsigned int secondPartLength = oddReadLength ? this->uiRead_Length - 1 : this->uiRead_Length;
    if (bClearQ) {
        aQue.clearHits();
    }
    if (!bForward) { //Query reverse complement read for alignment
        CReadInBits tmp = r2;
        r2 = reverseCompliment(firstPartLength, r1);
        r1 = reverseCompliment(firstPartLength, tmp);
    }
    CReadInBits readInBases = (queryHalf == 1) ? r1 : r2;
    bool bUseShortCut = this->bExcludeAmbiguous && !(aQue.qAllInThreshold());
    bool qAllHits = aQue.qAllInThreshold();
    for (unsigned int shift = 0; shift <= this->uiNoOfShift; shift++) {
        // If the best match is exact matched, no need to go to the next shift
        // for all exact matches must be found previously.
        if ((aQue.MinDiff == 0) && (shift > 0) && (aQue.load > 0) && !qAllHits) {
            break;
        }
        pair<CIndex_Type*, CIndex_Type*> hits = queryKmer(readInBases, shift);
        if (hits.first != NULL) {
            for (CIndex_Type* it = hits.first; it < hits.second; it++) {
                if (*it >= shift) {
                    unsigned int alignStart = *it - shift;
                    if (queryHalf == 2 ) {
                        if ( alignStart >= secondPartLength) {
                            alignStart -= secondPartLength;
                        } else {
                            continue;
                        }
                    }
                    if (isMasked(alignStart) || isMasked(alignStart + secondPartLength)) {
                        continue; // if the two windows contain N or is in references' border
                    } else {
                        unsigned int uiDiff = checkAlignment(alignStart, r1, r2, oddReadLength);
                        // The flag in alignmentsQ decide whether all alignments within uiDiff or
                        // only the alignments with minimum Diff are queue.
                        if (uiDiff <= this->uiSubDiffThreshold) {
                            aQue.saveHits(alignStart, (unsigned short)uiDiff);
                        }

                        if (bUseShortCut) { // short cut to exclude ambiguous reads
                            bool bMap2Repeat = this->pbaRepeatRepresentativeFlag->b(alignStart);
                            if (bMap2Repeat) {
                                aQue.AmbiguousFlag = true;
                            }
                            if (aQue.MinDiff == 0 && (aQue.load >= 2 || bMap2Repeat)) {
                                aQue.setForwardLoad(bForward);
                                return(0);
                            }
                        }
                    }
                }
            }
            if (aQue.MinDiff == 0) {
                // short cut. Output no more than iMaxCapacity exaxt alignment
                /* if (aQue.load >= aQue.iMaxCapacity - 1) {
                    aQue.setForwardLoad(bForward);
                    return(0);
                }*/
                // short cut. All exact matches will be found after first shift
                if (bUseShortCut && aQue.load >= 2) {
                    aQue.setForwardLoad(bForward);
                    return(1);
                }
            }
        }
    }
    // The records before are for the forward direction if this is a forward query
    aQue.setForwardLoad(bForward);
    return(aQue.MinDiff);
}

// TODO Complete the function
// Query a read in bases (Illumina) for hit and check uiDiff and put the result into the given Queue
unsigned int CGenome_Index_TableQ::queryLongReadColors(CReadInBits r1, CReadInBits r2, bool oddReadLength, CAlignmentsQ& aQue, int queryHalf, bool bClearQ, bool bForward) const
{
    // const unsigned int firstPartLength = this->uiRead_Length;
    const unsigned int secondPartLength = oddReadLength ? this->uiRead_Length - 1 : this->uiRead_Length;
    if (bClearQ) {
        aQue.clearHits();
    }
    if (!bForward) { // Query reverse complement read for alignment
        reverseLongColorRead(r1, r2, oddReadLength);
    }
    CReadInBits readInBases = (queryHalf == 1) ? r1 : r2;
    bool bUseShortCut = this->bExcludeAmbiguous && !(aQue.qAllInThreshold());
    bool qAllHits = aQue.qAllInThreshold();
    for (unsigned int shift = 0; shift <= this->uiNoOfShift; shift++) {
        // If the best match is exact matched, no need to go to the next shift
        // for all exact matches must be found previously.
        if ((aQue.MinDiff == 0) && (shift > 0) && (aQue.load > 0) && !qAllHits) {
            break;
        }
        pair<CIndex_Type*, CIndex_Type*> hits = queryKmer(readInBases, shift);
        if (hits.first != NULL) {
            for (CIndex_Type* it = hits.first; it < hits.second; it++) {
                if (*it >= shift) {
                    unsigned int alignStart = *it - shift;
                    if (queryHalf == 2 ) {
                        if ( alignStart >= secondPartLength) {
                            alignStart -= secondPartLength;
                        } else {
                            continue;
                        }
                    }
                    if (isMasked(alignStart) || isMasked(alignStart + secondPartLength)) {
                        continue; // if the two windows contain N or is in references' border
                    } else {
                        // TODO implement the checkColor
                        unsigned int uiDiff = checkColorAlignment(alignStart, r1, r2, oddReadLength);
                        // The flag in alignmentsQ decide whether all alignments within uiDiff or
                        // only the alignments with minimum Diff are queue.
                        if (uiDiff <= this->uiSubDiffThreshold) {
                            aQue.saveHits(alignStart, (unsigned short)uiDiff);
                        }
                        if (bUseShortCut) { // short cut to exclude ambiguous reads
                            bool bMap2Repeat = this->pbaRepeatRepresentativeFlag->b(alignStart);
                            if (bMap2Repeat) {
                                aQue.AmbiguousFlag = true;
                            }
                            if (aQue.MinDiff == 0 && (aQue.load >= 2 || bMap2Repeat)) {
                                aQue.setForwardLoad(bForward);
                                return(0);
                            }
                        }
                    }
                }
            }
            if (aQue.MinDiff == 0) {
                // short cut. All exact matches will be found after first shift
                if (bUseShortCut && aQue.load >= 2) {
                    aQue.setForwardLoad(bForward);
                    return(1);
                }
            }
        }
    }
    // The records before are for the forward direction if this is a forward query
    aQue.setForwardLoad(bForward);
    return(aQue.MinDiff);
}

// check match for long read
unsigned int CGenome_Index_TableQ::checkAlignment(unsigned int alignStartGenomeIndex, CReadInBits& half1, CReadInBits& half2, bool oddReadLength) const
{
    unsigned int uiDiff = 0;
    CReadInBits ref1, ref2;
    ref1 = this->pgenomeNTInBits->getSubstringInBits(alignStartGenomeIndex);
    uiDiff += bitsStrNCompare(ref1, half1, this->uiRead_Length);
    int secondPartStart;
    if (oddReadLength) {
        secondPartStart = alignStartGenomeIndex + this->uiRead_Length - 1;
        ref2 = this->pgenomeNTInBits->getSubstringInBits(secondPartStart);
        uiDiff += bitsStrMNCompare(ref2, half2, 1, this->uiRead_Length - 1); // skip the 1st (middle) base
    } else {
        secondPartStart = alignStartGenomeIndex + this->uiRead_Length;
        ref2 = this->pgenomeNTInBits->getSubstringInBits(secondPartStart);
        uiDiff += bitsStrNCompare(ref2, half2, this->uiRead_Length);
    }
    return(uiDiff);
}

// TODO complete the function for color alignment
unsigned int CGenome_Index_TableQ::checkColorAlignment(unsigned int alignStartGenomeIndex, CReadInBits& half1, CReadInBits& half2, bool oddReadLength) const
{
    unsigned int uiDiff = 0;
    CReadInBits ref1 = this->pgenomeNTInBits->getSubstringInBits(alignStartGenomeIndex);
    CReadInBits ref1InColors = bases2Colors(ref1);
    uiDiff += bitsStrNCompare(ref1InColors, half1, this->uiRead_Length);
    int secondPartStart = alignStartGenomeIndex + this->uiRead_Length - 1;
    if (oddReadLength) {
        CReadInBits ref2 = this->pgenomeNTInBits->getSubstringInBits(secondPartStart - 1);
        CReadInBits ref2InColor = bases2PureColors(ref2);
        uiDiff += bitsStrMNCompare(ref2, half2, 1, this->uiRead_Length - 1);
        // TODO Check the middle and last color signal.
        // Watch out the boundary
    } else {
        CReadInBits ref2 = this->pgenomeNTInBits->getSubstringInBits(secondPartStart - 1);
        CReadInBits ref2InColor = bases2PureColors(ref2);
        // Warning! The last base info is missing due to bases2PureColors
        uiDiff += bitsStrNCompare(ref2, half2, this->uiRead_Length - 1);
    }
    return(uiDiff);
}

// Given alignments in alignmentsQ, check reads can be also aligned in the extended position
unsigned int CGenome_Index_TableQ::extendAlignment\
(CAlignmentsQ& alignmentsQ, CReadInBits extendedReadHalf) const
{
    unsigned int i;
    int minDiff = this->uiRead_Length;
    bool extendForward = true;
    for (i = 0; i < alignmentsQ.ForwardAlignmentLoad; i++) {
        unsigned int alignStart = alignmentsQ.aiHitIndex[i] + this->uiRead_Length;
        if (alignStart + this->uiRead_Length >= this->pgenomeNT->iGenomeSize || this->isMasked(alignStart)) {
            alignmentsQ.asdiff[i] += (unsigned short)this->uiRead_Length; // no space to align
        } else {
            CReadInBits ref = this->pgenomeNTInBits->getSubstringInBits(alignStart, this->uiRead_Length);
            alignmentsQ.asdiff[i] += (unsigned short)bitsStrNCompare(ref, extendedReadHalf, this->uiRead_Length);
            if (alignmentsQ.asdiff[i] < minDiff)
                minDiff = alignmentsQ.asdiff[i];
        }
    }
    extendForward = false;
    CReadInBits reverseComplimentRead = reverseCompliment(this->uiRead_Length, extendedReadHalf);
    for (; i < alignmentsQ.load; i++) { // reverse read
        if ( alignmentsQ.aiHitIndex[i] <  this->uiRead_Length) {
            alignmentsQ.asdiff[i] += (unsigned short)this->uiRead_Length; // no space to align
        } else {
            unsigned int alignStart = alignmentsQ.aiHitIndex[i] - this->uiRead_Length;
            if ( this->isMasked(alignStart) ) {
                alignmentsQ.asdiff[i] += (unsigned short)this->uiRead_Length; // no space to align
            } else {
                CReadInBits ref = this->pgenomeNTInBits->getSubstringInBits(alignStart, this->uiRead_Length);
                alignmentsQ.asdiff[i] += (unsigned short)bitsStrNCompare(ref, reverseComplimentRead, this->uiRead_Length);
                alignmentsQ.aiHitIndex[i] = alignStart;
                if (alignmentsQ.asdiff[i] < minDiff) {
                    minDiff = alignmentsQ.asdiff[i];
                }
            }
        }
    }
    alignmentsQ.MinDiff = minDiff;
    // The alignments may exceed the threshold so it should be checked outside the function.
    return(alignmentsQ.MinDiff);
}

// Query a read in colors (SOLiD) for hit and check uiDiff and put the result into the given Queue
unsigned int CGenome_Index_TableQ::queryReadColors(CReadInBits readInColors, CAlignmentsQ& alignmentsQ, bool bClearQ, bool bForward, bool bDEBUG) const
{
    if (bClearQ) {
        alignmentsQ.clearHits();
    }

    CReadInBits pureColors = readInColors.getSuffixStr(1); // query colors
    if (!bForward) { // Query alignment with reversed color. (Reverse complement reads has color in reversed direction)
        pureColors = reversePureColors(pureColors, this->uiRead_Length - 1);
    }

    bool bUseShortCut = this->bExcludeAmbiguous && !(alignmentsQ.qAllInThreshold());
    bool qAllHits = alignmentsQ.qAllInThreshold();
    for (unsigned int shift = 0; shift <= this->uiNoOfShift; shift++) {
        // If the best alignment is exact matched, no need to go to the next shift, for all exact matches must be found previously.
        if ((alignmentsQ.MinDiff == 0) && (shift > 0) && (alignmentsQ.load > 0) && !qAllHits) {
            break;
        }
        /*
        if(bDEBUG) {
        	cout << shift << "-Shift" << endl;
        	if(shift == 7) {
        		cout << "Got you" << endl;
        	}
        }*/

        // printBitsStr(pureColors.getSuffixStr(shift), this->uiSeedLength);// DEBUG
        pair<CIndex_Type*, CIndex_Type*> hits = queryKmer(pureColors, shift);
        if (hits.first != NULL) {  // for each hit, check if it is a good alignment.
            for (CIndex_Type* it = hits.first; it < hits.second; it++) {
                if (*it >= shift) {
                    unsigned int alignStart = *it - shift;
                    if (isMasked(alignStart)) {
                        continue;
                    } else {
                        CReadInBits ref = this->pgenomeNTInBits->getSubstringInBits(alignStart, this->uiRead_Length);
                        unsigned int uiDiff;
                        if (bForward) {
                            CReadInBits refInColors = bases2Colors(ref);
                            uiDiff = bitsStrNCompare(refInColors, readInColors, this->uiRead_Length);
                        } else { //reverse complement the reference to compare in the read
                            CReadInBits rcRefInColors = bases2Colors(reverseCompliment(this->uiRead_Length, ref));
                            uiDiff = bitsStrNCompare(rcRefInColors, readInColors, this->uiRead_Length);
                        }
                        // To queue all alignment within uiDiff or the alignments with min difference are set in flag of alignmentsQ
                        if (uiDiff <= this->uiSubDiffThreshold) {
                            alignmentsQ.saveHits(alignStart, (unsigned short)uiDiff);
                        }

                        if (bUseShortCut) { // short cut to exclude ambiguous reads
                            bool bMap2Repeat = this->pbaRepeatRepresentativeFlag->b(alignStart);
                            if (bMap2Repeat) {
                                alignmentsQ.AmbiguousFlag = true;
                            }
                            if (alignmentsQ.MinDiff == 0 && (alignmentsQ.load >= 2 || bMap2Repeat)) {
                                alignmentsQ.setForwardLoad(bForward);
                                return(0);
                            }
                        }
                    }
                }
            }
            if (alignmentsQ.MinDiff == 0) {
                /*
                // short cut to exclude reads that has too many ambiguous mapping
                if (alignmentsQ.load >= MAX_Q_CAPACITY) {
                    alignmentsQ.setForwardLoad(bForward);
                    return(alignmentsQ.MinDiff);
                }
                */
                // short cut for ambiguous reads
                if (bUseShortCut && alignmentsQ.load >= 2) {
                    alignmentsQ.setForwardLoad(bForward);
                    return(1);
                }
            }
        }
    }
    alignmentsQ.setForwardLoad(bForward);
    return(alignmentsQ.MinDiff);
}

