#ifndef _BLASR_AFFINE_KBAND_ALIGN_HPP_
#define _BLASR_AFFINE_KBAND_ALIGN_HPP_

#include <cassert>
#include <iostream>
#include <vector>

#include <pbdata/defs.h>
#include <alignment/algorithms/alignment/KBandAlign.hpp>
#include <alignment/datastructures/alignment/Alignment.hpp>
#include <pbdata/NucConversion.hpp>
#include <pbdata/matrix/FlatMatrix.hpp>

template <typename T_QuerySequence, typename T_TargetSequence, typename T_Alignment>
int AffineKBandAlign(T_QuerySequence &pqSeq, T_TargetSequence &ptSeq, int matchMat[5][5],
                     int hpInsOpen, int hpInsExtend, int insOpen, int insExtend, int del, int k,
                     std::vector<int> &scoreMat, std::vector<Arrow> &pathMat,
                     std::vector<int> &hpInsScoreMat, std::vector<Arrow> &hpInsPathMat,
                     std::vector<int> &insScoreMat, std::vector<Arrow> &insPathMat,
                     T_Alignment &alignment, AlignmentType alignType)
{

    //
    // Make a copy of the sequences that is guaranteed to be in 3-bit format
    // for quick access to the score array.
    //

    int INF_SCORE = INF_INT - 1000;
    T_QuerySequence qSeq;
    T_TargetSequence tSeq;
    //	CreateThreeBitSequence(pqSeq, qSeq);
    //	CreateThreeBitSequence(ptSeq, tSeq);
    qSeq.seq = pqSeq.seq;
    qSeq.length = pqSeq.length;
    tSeq.seq = ptSeq.seq;
    tSeq.length = ptSeq.length;

    DNALength tLen, qLen;
    SetKBoundedLengths(tSeq.length, qSeq.length, k, tLen, qLen);

    //
    //
    // Allow for length up to diagonal + k + 1 for boundary.
    //
    // Allow for width:
    //   diagonal (1)
    //   up to k insertions (k)
    //   up to k deletions  (k)
    //   boundary on left side of matrix (1)
    //
    //	if (qLen + k > tLen and qLen < tLen) {
    //		k = tLen - qLen +1;
    //	}
    DNALength nCols = 2 * k + 1;
    VectorIndex totalMatSize = (qLen + 1) * nCols;

    //
    // For now the scoreMat and path mat maintained outside this
    // function so that they may have different sizes from the affine
    // matrices.
    //

    if (scoreMat.size() < totalMatSize) {
        scoreMat.resize(totalMatSize);
        pathMat.resize(totalMatSize);
    }

    if (hpInsScoreMat.size() < totalMatSize) {
        hpInsScoreMat.resize(totalMatSize);
        hpInsPathMat.resize(totalMatSize);
        insScoreMat.resize(totalMatSize);
        insPathMat.resize(totalMatSize);
    }

    //
    // Initialze matrices
    //
    std::fill(scoreMat.begin(), scoreMat.begin() + totalMatSize, 0);
    std::fill(pathMat.begin(), pathMat.begin() + totalMatSize, NoArrow);
    std::fill(hpInsScoreMat.begin(), hpInsScoreMat.begin() + totalMatSize, 0);
    std::fill(hpInsPathMat.begin(), hpInsPathMat.begin() + totalMatSize, NoArrow);
    std::fill(insScoreMat.begin(), insScoreMat.begin() + totalMatSize, 0);
    std::fill(insPathMat.begin(), insPathMat.begin() + totalMatSize, NoArrow);

    //
    // Initialize the boundaries of the DP matrix.
    //
    int q, t;

    if (alignType != TargetFit) {
        insScoreMat[rc2index(0, k, nCols)] = 0;
        insPathMat[rc2index(0, k, nCols)] = AffineInsOpen;

        for (q = 1; q <= k && q < static_cast<int>(qLen) + 1; q++) {
            insScoreMat[rc2index(q, k - q, nCols)] = q * insExtend + insOpen;
            insPathMat[rc2index(q, k - q, nCols)] = AffineInsUp;
        }
    } else if (alignType == TargetFit) {
        //
        // Allow free gap penalties at the beginning of the alignment.
        //
        insScoreMat[rc2index(0, k, nCols)] = 0;
        insPathMat[rc2index(0, k, nCols)] = AffineInsOpen;
        for (q = 1; q <= k && q < static_cast<int>(qLen) + 1; q++) {
            insScoreMat[rc2index(q, k - q, nCols)] = 0;
            insPathMat[rc2index(q, k - q, nCols)] = AffineInsUp;
        }
    }

    //
    // Assign score for (0,0) position in matrix -- aligning a gap to a gap
    // which should just be a finished alignment.  There is no cost for
    // gap-gap alignment.
    //
    hpInsScoreMat[rc2index(0, k, nCols)] = 0;
    hpInsPathMat[rc2index(0, k, nCols)] = AffineHPInsOpen;

    for (q = 1; q <= k && q < static_cast<int>(qLen) + 1; q++) {
        hpInsScoreMat[rc2index(q, k - q, nCols)] = q * hpInsExtend + hpInsOpen;
        hpInsPathMat[rc2index(q, k - q, nCols)] = AffineHPInsUp;
    }

    for (t = k + 1; t < static_cast<int>(nCols); t++) {
        hpInsScoreMat[rc2index(0, t, nCols)] =
            INF_SCORE;               //  hpInsOpen + (t - k) * hpInsExtend;//; //INF_SCORE;
        hpInsPathMat[t] = NoArrow;   //AffineHPInsOpen ; //NoArrow;
        insScoreMat[t] = INF_SCORE;  //insOpen + (t - k) * insExtend; //INF_SCORE;
        insPathMat[t] = NoArrow;     //AffineInsOpen; //NoArrow;
    }

    for (q = 1; q <= k && q < static_cast<int>(qLen) + 1; q++) {
        scoreMat[rc2index(q, k - q, nCols)] = insScoreMat[rc2index(q, k - q, nCols)];
        pathMat[rc2index(q, k - q, nCols)] = AffineInsClose;
    }
    for (t = 1; t <= k; t++) {
        scoreMat[rc2index(0, t + k, nCols)] = t * del;
        pathMat[rc2index(0, t + k, nCols)] = Left;
    }

    //
    // The recurrence relation here is a slight modification of the
    // standard affine gap alignment.  Deletions are non-affine.  Insertions
    // are affine with different scores for homopolymer insertions, and
    // an affine score for mixed insertions.
    //

    int matchScore, delScore;
    int hpInsExtendScore, hpInsOpenScore, insOpenScore, insExtendScore;
    int minHpInsScore, minInsScore;
    for (q = 1; q <= static_cast<int>(qLen); q++) {
        for (t = q - k; t < static_cast<int>(q) + k + 1; t++) {
            if (t < 1) {
                continue;
            }
            if (static_cast<DNALength>(t) > tLen) {
                break;
            }

            VectorIndex upper = rc2index(q - 1, k + t - q + 1, nCols);
            VectorIndex curIndex = rc2index(q, k + t - q, nCols);

            if (t < q + k)
                hpInsOpenScore = scoreMat[upper] + hpInsOpen;
            else
                hpInsOpenScore = INF_SCORE;

            //
            // The homopolymer insertion score is defined only when the previous nucleotide
            // is the same as the current, in which case the homopolymer insertion score
            // is used.  If the current and previous nucleotide in the query are different,
            // the extension is not possible, and the best that can happen is a gap open.
            //
            if (q > 1 and qSeq[q - 1] == qSeq[q - 2]) {
                if (t < q + k)
                    hpInsExtendScore = hpInsScoreMat[upper] + hpInsExtend;
                else
                    hpInsExtendScore = INF_SCORE;
            } else {
                hpInsExtendScore = INF_SCORE;
            }

            //
            // Since this is only allowing insertions, this grid has only horizontal and
            // elevation arrows.
            //

            if (hpInsOpenScore < hpInsExtendScore) {
                hpInsPathMat[curIndex] = AffineHPInsOpen;
                minHpInsScore = hpInsOpenScore;
            } else {
                hpInsPathMat[curIndex] = AffineHPInsUp;
                minHpInsScore = hpInsExtendScore;
            }

            hpInsScoreMat[curIndex] = minHpInsScore;
            if (t < q + k) {
                insOpenScore = scoreMat[upper] + insOpen;
                insExtendScore = insScoreMat[upper] + insExtend;
            } else {
                insOpenScore = INF_SCORE;
                insExtendScore = INF_SCORE;
            }

            if (insOpenScore < insExtendScore) {
                insPathMat[curIndex] = AffineInsOpen;
                minInsScore = insOpenScore;
            } else {
                insPathMat[curIndex] = AffineInsUp;
                minInsScore = insExtendScore;
            }
            insScoreMat[curIndex] = minInsScore;

            // On left boundary of k-band.
            // do not allow deletions of t.
            if (t == q - k) {
                delScore = INF_SCORE;
            } else {
                // cur row = q
                // cur col = t - q
                // prev col therefore t - q - 1
                // and offset from diagonal is k + t - q - 1
                delScore = scoreMat[rc2index(q, k + t - q - 1, nCols)] + del;
            }

            // cur row = q
            // cur col = t - q

            // cur query index = q - 1
            // cur target index = t - 1
            // therefore match row (up) = q
            //           match col (left, but since up shifted right) = t - q
            assert(rc2index(q - 1, k + t - q, nCols) < scoreMat.size());
            assert(t - 1 >= 0);
            assert(q - 1 >= 0);
            matchScore = scoreMat[rc2index(q - 1, k + t - q, nCols)] +
                         matchMat[ThreeBit[qSeq.seq[q - 1]]][ThreeBit[tSeq.seq[t - 1]]];

            //
            //  Possibly on right boundary of k-band, in which
            //  case do not allow insertions from q.

            int minScore = MIN(matchScore, MIN(delScore, MIN(minInsScore, minHpInsScore)));
            curIndex = rc2index(q, k + t - q, nCols);
            assert(curIndex < scoreMat.size());
            scoreMat[curIndex] = minScore;
            if (minScore == matchScore) {
                pathMat[curIndex] = Diagonal;
            } else if (minScore == delScore) {
                pathMat[curIndex] = Left;
            } else if (minScore == minInsScore) {
                pathMat[curIndex] = AffineInsClose;
            } else {
                pathMat[curIndex] = AffineHPInsClose;
            }
        }
    }
    /*
       std::cout << "tracing back from: " << q << ", " << t << std::endl;
       std::cout << "match score: " << std::endl;
       PrintFlatMatrix(&scoreMat[0], qLen + 1, nCols, std::cout);
       std::cout << " path: " << std::endl;
       PrintFlatMatrix(pathMat, qLen + 1, nCols, std::cout);
       std::cout << "hp  score: " << std::endl;
       PrintFlatMatrix(&hpInsScoreMat[0], qLen + 1, nCols, std::cout);
       std::cout << "hp  path: " << std::endl;
       PrintFlatMatrix(hpInsPathMat, qLen + 1, nCols, std::cout);
       std::cout << "normal affine ins score: " << std::endl;
       PrintFlatMatrix(&insScoreMat[0], qLen + 1, nCols, std::cout);
       std::cout << "normal affine ins path: " << std::endl;
       PrintFlatMatrix(&insPathMat[0], qLen + 1, nCols, std::cout);
       */
    std::vector<Arrow> optAlignment;
    // First find the end position matrix.

    int minScoreTPos, minScore;
    int minScoreQPos;
    if (alignType == Global) {
        q = qLen;
        t = k - (static_cast<int>(qLen) - static_cast<int>(tLen));
    } else if (alignType == QueryFit) {
        q = qLen;
        minScoreTPos = std::max(q - k, 1);
        DNALength index = rc2index(qLen, k + minScoreTPos - q, nCols);
        minScore = scoreMat[index];
        for (t = q - k; t < q + k + 1; t++) {
            if (t < 1) {
                continue;
            }
            if (t > static_cast<int>(tLen)) {
                break;
            }
            int index = rc2index(qLen, k + t - q, nCols);
            if (scoreMat[index] < minScore) {
                minScoreTPos = t;
                minScore = scoreMat[index];
            }
        }
        t = k - (static_cast<int>(qLen) - minScoreTPos);
    } else if (alignType == TargetFit) {
        t = tLen;

        int qStart = std::max(0, std::min((int)qLen, (int)tLen) -
                                     std::max(0, k - std::max(((int)tLen) - ((int)qLen), 0)));
        int qEnd = std::min(qLen, tLen + k) + 1;

        minScoreQPos = qStart;
        int index = rc2index(minScoreQPos, k - (minScoreQPos - tLen), nCols);
        minScore = scoreMat[index];
        for (q = qStart; q < qEnd; q++) {
            // add to k since this is going up.
            index = rc2index(q, k + (q - tLen), nCols);
            if (scoreMat[index] < minScore) {
                minScoreQPos = q;
                minScore = scoreMat[index];
            }
        }
        q = minScoreQPos;
        t = (k + ((int)q - (int)tLen));
    }

    int optScore = scoreMat[rc2index(q, t, nCols)];
    Arrow arrow;
    MatrixLabel curMatrix = Match;

    while ((q > 0) or (q == 0 and t > k)) {
        assert(t < 2 * k + 1);
        if (curMatrix == Match) {
            arrow = pathMat[rc2index(q, t, nCols)];
            if (arrow == Diagonal) {
                optAlignment.push_back(arrow);
                q--;
            } else if (arrow == Left) {
                optAlignment.push_back(arrow);
                t--;
            }
            //
            // The following two conditions change matrices
            // without changing coordinates, since the gap close
            // just changes state without adding to the alignment.
            //
            else if (arrow == AffineInsClose) {
                curMatrix = AffineIns;
            } else if (arrow == AffineHPInsClose) {
                curMatrix = AffineHPIns;
            }
        } else if (curMatrix == AffineHPIns) {
            //
            // The current
            arrow = hpInsPathMat[rc2index(q, t, nCols)];
            if (arrow == AffineHPInsOpen) {
                curMatrix = Match;
            } else if (arrow != AffineHPInsUp) {
                std::cout << "ERROR! Affine homopolymer insertion path matrix MUST only have UP or "
                             "OPEN arrows."
                          << std::endl;
                assert(0);
            }
            optAlignment.push_back(Up);
            q--;
            t++;
        } else if (curMatrix == AffineIns) {
            arrow = insPathMat[rc2index(q, t, nCols)];
            if (arrow == AffineInsOpen) {
                curMatrix = Match;
            } else if (arrow != AffineInsUp) {
                std::cout << "ERROR! Affine insertion path matrix MUST only have UP or OPEN arrows."
                          << std::endl;
                assert(0);
            }
            optAlignment.push_back(Up);
            q--;
            t++;
        } else {
            std::cout << "ERROR in affine local alignment, matrix is: " << curMatrix << std::endl;
            assert(0);
        }
    }
    //	qSeq.Free();
    //	tSeq.Free();
    std::reverse(optAlignment.begin(), optAlignment.end());
    alignment.ArrowPathToAlignment(optAlignment);
    return optScore;
}

#endif  // _BLASR_AFFINE_KBAND_ALIGN_HPP_
