/* bandExt - banded Smith-Waterman extension of alignments. 
 * An aligner might first find perfectly matching hits of
 * a small size, then extend these hits as far as possible
 * while the sequences perfectly match, then call on routines
 * in this module to do further extensions allowing small
 * gaps and mismatches. 
 *
 * This version of the algorithm uses affine gap scorea and
 * has the neat feature that the band can wander around.  
 * When a score exceeds any previous score, the band will be 
 * recentered around the highest scoring position. */
/* Copyright 2003-5 Jim Kent.  All rights reserved. */

#include "common.h"
#include "dnaseq.h"
#include "axt.h"
#include "fuzzyFind.h"
#include "localmem.h"
#include "bandExt.h"


/* Definitions for traceback byte.  This is encoded as so:
 *     xxxxLUMM
 * where the x's are not used.  The L bit is set if parent of the left insert
 * is also a left insert (otherwise it is a match). Likewise the U bit is set
 * if the parent of an up insert is also an up insert.  The MM bits contain the
 * parent of the match state. */

#define mpMatch 1
#define mpUp 2
#define mpLeft 3
#define mpMask 3
#define upExt (1<<2)
#define lpExt (1<<3)

struct score
/* Info on score in our three states. */
   {
   int match;
   int up;
   int left;
   };

boolean bandExt(boolean global, struct axtScoreScheme *ss, int maxInsert,
	char *aStart, int aSize, char *bStart, int bSize, int dir,
	int symAlloc, int *retSymCount, char *retSymA, char *retSymB, 
	int *retStartA, int *retStartB)
/* Try to extend an alignment from aStart/bStart onwards.
 * Set maxInsert to the maximum gap size allowed.  3 is often
 * a good choice.  aStart/aSize bStart/bSize describe the
 * sequence to extend through (not including any of the existing
 * alignment. Set dir = 1 for forward extension, dir = -1 for backwards. 
 * retSymA and retSymB should point to arrays of characters of
 * symAlloc size.  symAlloc needs to be aSize*2 or so.  The
 * alignment is returned in the various ret values.  The function
 * overall returns TRUE if an extension occurred, FALSE if not. */
{
int i;			/* A humble index or two. */
int *bOffsets = NULL;	/* Offset of top of band. */
UBYTE **parents = NULL;	/* Array of parent positions. */
struct score *curScores = NULL;	/* Scores for current column. */
struct score *prevScores = NULL;	/* Scores for previous column. */
struct score *swapScores;	/* Helps to swap cur & prev column. */
int bandSize = 2*maxInsert + 1;	 /* Size of band including middle */
int maxIns1 = maxInsert + 1;   	 /* Max insert plus one. */
int bandPlus = bandSize + 2*maxIns1;  /* Band plus sentinels on either end. */
int bestScore = 0;	 	 /* Best score so far. */
int aPos, aBestPos = -1;	 /* Current and best position in a. */
int bPos, bBestPos = -1;	 /* Current and best position in b. */
int bandCenter = 0;	/* b Coordinate of current band center. */
int colShift = 1;	/* Vertical shift between this column and previous. */
int prevScoreOffset;	/* Offset into previous score column. */
int curScoreOffset;	/* Offset into current score column. */
int symCount = 0;	/* Size of alignment and size allocated for it. */
int gapOpen = ss->gapOpen;	 /* First base in gap costs this. */
int gapExtend = ss->gapExtend;	 /* Extra bases in gap cost this. */
int badScore = -gapOpen*100;	 /* A score bad enough no one will want to link with us. */
int maxDrop = gapOpen + gapExtend*maxInsert; /* Max score drop allowed before giving up. */
int midScoreOff;		 /* Offset to middle of scoring array. */
struct lm *lm;			 /* Local memory pool. */
boolean didExt = FALSE;
int initGapScore = -gapOpen;

/* For reverse direction just reverse bytes here and there.  It's
 * a lot easier than the alternative and doesn't cost much time in
 * the global scheme of things. */
if (dir < 0)
    {
    reverseBytes(aStart, aSize);
    reverseBytes(bStart, bSize);
    }
#ifdef DEBUG
uglyf("bandExt: dir %d, aStart %d, aSize %d, symAlloc %d\n", dir,
	aSize, bSize, symAlloc);
mustWrite(uglyOut, aStart, aSize);
uglyf("\n");
mustWrite(uglyOut, bStart, bSize);
uglyf("\n");
#endif /* DEBUG */

/* Make up a local mem structure big enough for everything. 
 * This is just a minor, perhaps misguided speed tweak to
 * avoid multiple mallocs. */
lm = lmInit(
    sizeof(bOffsets[0])*aSize +
    sizeof(parents[0])*bandSize +
    bandSize*(sizeof(parents[0][0])*aSize) +
    sizeof(curScores[0]) * bandPlus +
    sizeof(prevScores[0]) * bandPlus );

/* Allocate data structures out of local memory pool. */
lmAllocArray(lm, bOffsets, aSize);
lmAllocArray(lm, parents, bandSize);
for (i=0; i<bandSize; ++i)
    lmAllocArray(lm, parents[i], aSize);
lmAllocArray(lm, curScores, bandPlus);
lmAllocArray(lm, prevScores, bandPlus);

/* Set up scoring arrays so that stuff outside of the band boundary 
 * looks bad.  There will be maxIns+1 of these sentinel values at
 * the start and at the beginning. */
for (i=0; i<bandPlus; ++i)
    {
    struct score *cs = &curScores[i];
    cs->match = cs->up = cs->left = badScore;
    cs = &prevScores[i];
    cs->match = cs->up = cs->left = badScore;
    }

/* Set up scoring array so that extending without an initial insert
 * looks relatively good. */
midScoreOff = 1 + 2 * maxInsert;
prevScores[midScoreOff].match = 0;

/* Set up scoring array so that initial inserts of up to maxInsert
 * are allowed but penalized appropriately. */
    {
    int score = -gapOpen;
    for (i=0; i<maxInsert; ++i)
	{
	prevScores[midScoreOff+i].up = score;
	score -= gapExtend;
	}
    }

for (aPos=0; aPos < aSize; ++aPos)
    {
    char aBase = aStart[aPos];
    int *matRow = ss->matrix[(int)aBase];
    int bestColScore = badScore;
    int bestColPos = -1;
    int colTop = bandCenter - maxInsert;
    int colBottom = bandCenter + maxIns1;

    if (colTop < 0) colTop = 0;
    if (colBottom > bSize) colBottom = bSize;
    curScoreOffset =  maxIns1 + colTop - (bandCenter - maxInsert);
    prevScoreOffset = curScoreOffset + colShift;
#ifdef DEBUG
    uglyf("curScoreOffset %d, maxIns %d, prevScoreOffset %d\n", curScoreOffset, maxInsert, prevScoreOffset);
#endif /* DEBUG */

    /* At top we deal with possibility of initial insert that
     * comes in at this point, unless the band has wandered off. */
    if (aPos < maxInsert)
	{
	curScores[curScoreOffset-1].up = initGapScore;
	initGapScore -= gapExtend;
	}
    else
	curScores[curScoreOffset-1].up = badScore;

    for (bPos = colTop; bPos < colBottom; ++bPos)
	{
	UBYTE parent;
	struct score *curScore = &curScores[curScoreOffset];

#ifdef DEBUG
	uglyf("aPos %d, bPos %d, %c vs %c\n", aPos, bPos, aBase, bStart[bPos]);
	uglyf("  diag[%d %d %d], left[%d %d %d], up[%d %d %d]\n", 
		prevScores[prevScoreOffset-1].match, 
		prevScores[prevScoreOffset-1].left, 
		prevScores[prevScoreOffset-1].up, 
		prevScores[prevScoreOffset].match, 
		prevScores[prevScoreOffset].left, 
		prevScores[prevScoreOffset].up, 
		curScores[curScoreOffset-1].match, 
		curScores[curScoreOffset-1].left, 
		curScores[curScoreOffset-1].up);
#endif /* DEBUG */

	/* Handle ways into the matching state and record if it's
	 * best score so far. */
	    {
	    int match = matRow[(int)bStart[bPos]];
	    struct score *a = &prevScores[prevScoreOffset-1];
	    int diagScore = a->match;
	    int leftScore = a->left;
	    int upScore = a->up;
	    int score;
	    if (diagScore >= leftScore && diagScore >= upScore)
	        {
		score = diagScore + match;
		parent = mpMatch;
		}
	    else if (leftScore > upScore)
	        {
		score = leftScore + match;
	        parent = mpLeft;
		}
	    else
	        {
	        score = upScore + match;
	        parent = mpUp;
	        }
	    curScore->match = score;
	    if (score > bestColScore)
	        {
		bestColScore = score;
		bestColPos = bPos;
		}
	    }

	/* Handle ways into left gap state. */
	    {
	    struct score *a = &prevScores[prevScoreOffset];
	    int extScore = a->left - gapExtend;
	    int openScore = a->match - gapOpen;
	    if (extScore >= openScore)
	        {
		parent |= lpExt;
		curScore->left = extScore;
		}
	    else
	        {
		curScore->left = openScore;
		}
	    }

	/* Handle ways into up gap state. */
	    {
	    struct score *a = &curScore[-1];
	    int extScore = a->up - gapExtend;
	    int openScore = a->match - gapOpen;
	    if (extScore >= openScore)
	        {
		parent |= upExt;
		curScore->up = extScore;
		}
	    else
	        {
		curScore->up = openScore;
		}
	    }

	parents[curScoreOffset - maxIns1][aPos] = parent;

#ifdef DEBUG
	uglyf(" cur [%d %d %d]\n", 
		curScore->match, 
		curScore->left, 
		curScore->up);
#endif /* DEBUG */
	/* Advance to next row in column. */
	curScoreOffset += 1;
	prevScoreOffset += 1;
	}


    /* If this column's score is best so far make note of
     * it and shift things so that the matching bases at the
     * best score are centered in the band.  This allows the
     * band to wander where appropriate.  Having the band shift
     * to the best column position if it is not the best score
     * so far doesn't work so well though.  */
    if (bestScore < bestColScore)
	{
	bestScore = bestColScore;
	aBestPos = aPos;
	bBestPos = bestColPos;
	colShift = (bestColPos + 1) - bandCenter;
	}
    else if (bestColScore < bestScore - maxDrop)
	{
	if (!global)
	    break;
	}
    else
	{
	colShift = 1;
	}

    /* Keep track of vertical offset of this column, and
     * set up offset of next column. */
    bOffsets[aPos] = bandCenter;
    bandCenter += colShift;

    /* Swap scoring arrays so current becomes next iteration's previous. */
    swapScores = curScores;
    curScores = prevScores;
    prevScores = swapScores;
    }


/* Trace back. */
#ifdef DEBUG
uglyf("traceback: bestScore %d, aBestPos %d, bBestPos %d\n", bestScore, aBestPos, bBestPos);
#endif /* DEBUG */
if (global || bestScore > 0)
    {
    boolean upState = FALSE;
    boolean leftState = FALSE;
    didExt = TRUE;
    if (global)
        {
	aPos = aSize-1;
	bPos = bSize-1;
	}
    else
	{
	aPos = aBestPos;
	bPos = bBestPos;
	}
    for (;;)
	{
	int pOffset;
	UBYTE parent;
	pOffset = bPos - bOffsets[aPos] + maxInsert;
	if (pOffset < 0) pOffset = 0;
	// FIXME: there may be some cases near beginning where
	// pOffset is not right.  Clamping it at 0 prevents a crash
	// and produces correct behavior in many cases.  However
	// I'm thinking a better fix would be to have bOffsets follow
	// colTop rather than bandCenter somehow. Exactly how is
	// beyond me at the moment.
	if (pOffset >= bandSize)
	    {
#ifdef DEBUG
	    uglyf("bPos %d, aPos %d, bOffsets[aPos] %d, maxInsert %d\n", bPos, aPos, bOffsets[aPos], maxInsert);
	    uglyf("pOffset = %d, bandSize = %d\n", pOffset, bandSize);
	    uglyf("aSize %d, bSize %d\n", aSize, bSize);
	    faWriteNext(uglyOut, "uglyA", aStart, aSize);
	    faWriteNext(uglyOut, "uglyB", bStart, bSize);
#endif /* DEBUG */
	    assert(global);
	    return FALSE;
	    }
	parent = parents[pOffset][aPos];
#ifdef DEBUG
	uglyf("aPos %d, bPos %d, parent %d, pMask %d upState %d, leftState %d\n", aPos, bPos, parent, parent & mpMask, upState, leftState);
#endif /* DEBUG */
	if (upState)
	    {
	    retSymA[symCount] = '-';
	    retSymB[symCount] = bStart[bPos];
	    bPos -= 1;
	    upState = (parent&upExt);
	    }
	else if (leftState)
	    {
	    retSymA[symCount] = aStart[aPos];
	    retSymB[symCount] = '-';
	    aPos -= 1;
	    leftState = (parent&lpExt);
	    }
	else
	    {
	    retSymA[symCount] = aStart[aPos];
	    retSymB[symCount] = bStart[bPos];
	    aPos -= 1;
	    bPos -= 1;
	    parent &= mpMask;
	    if (parent == mpUp)
	        upState = TRUE;
	    else if (parent == mpLeft)
	        leftState = TRUE;
	    }
	if (++symCount >= symAlloc)
	    errAbort("unsufficient symAlloc in bandExt");
	if (aPos < 0 || bPos < 0)
	    {
	    while (aPos >= 0)
		{
		retSymA[symCount] = aStart[aPos];
		retSymB[symCount] = '-';
		if (++symCount >= symAlloc)
		    errAbort("unsufficient symAlloc in bandExt");
		--aPos;
		}
	    while (bPos >= 0)
		{
		retSymA[symCount] = '-';
		retSymB[symCount] = bStart[bPos];
		if (++symCount >= symAlloc)
		    errAbort("unsufficient symAlloc in bandExt");
		--bPos;
		}
	    break;
	    }
	}
    if (dir > 0)
	{
	reverseBytes(retSymA, symCount);
	reverseBytes(retSymB, symCount);
	}
    retSymA[symCount] = 0;
    retSymB[symCount] = 0;
    }
else
    {
    retSymA[0] = retSymB[0] = 0;
    }

if (dir < 0)
    {
    reverseBytes(aStart, aSize);
    reverseBytes(bStart, bSize);
    }

/* Clean up, set return values and go home */
lmCleanup(&lm);
if (retStartA != NULL) *retStartA = aBestPos;
if (retStartB != NULL) *retStartB = bBestPos;
*retSymCount = symCount;
return didExt;
}

struct ffAli *bandExtFf(
	struct lm *lm,	/* Local memory pool, NULL to use global allocation for ff */
	struct axtScoreScheme *ss, 	/* Scoring scheme to use. */
	int maxInsert,			/* Maximum number of inserts to allow. */
	struct ffAli *origFf,		/* Alignment block to extend. */
	char *nStart, char *nEnd,	/* Bounds of region to extend through */
	char *hStart, char *hEnd,	/* Bounds of region to extend through */
	int dir,			/* +1 to extend end, -1 to extend start */
	int maxExt)			/* Maximum length of extension. */
/* Extend a gapless alignment in one direction.  Returns extending
 * ffAlis, not linked into origFf, or NULL if no extension possible. */
{
int symAlloc = 2*maxExt;
char *symBuf = needMem(4*maxExt);
char *nBuf = symBuf;
char *hBuf = symBuf + symAlloc;
char *ns, *hs;
int symCount, nExt, hExt;
int nSize, hSize;
struct ffAli *ffList = NULL;
boolean gotExt;

if (dir > 0)
    {
    nSize = nEnd - origFf->nEnd;
    hSize = hEnd - origFf->hEnd;
    if (nSize > maxExt) nSize = maxExt;
    if (hSize > maxExt) hSize = maxExt;
    ns = origFf->nEnd;
    hs = origFf->hEnd;
    }
else
    {
    nSize = origFf->nStart - nStart;
    hSize = origFf->hStart - hStart;
    if (nSize > maxExt) nSize = maxExt;
    if (hSize > maxExt) hSize = maxExt;
    ns = origFf->nStart - nSize;
    hs = origFf->hStart - hSize;
    }

gotExt = bandExt(FALSE, ss, maxInsert, ns, nSize, hs, hSize, dir,
	symAlloc, &symCount, nBuf, hBuf, &nExt, &hExt);
if (gotExt)
    {
    char *nExtStart, *hExtStart;
    if (dir > 0)
	{
	nExtStart = ns;
	hExtStart = hs;
	}
    else
	{
	nExtStart = origFf->nStart - nExt - 1;
	hExtStart = origFf->hStart - hExt - 1;
	}
    ffList = ffAliFromSym(symCount, nBuf, hBuf, lm,  nExtStart, hExtStart);
    }
freeMem(symBuf);
return ffList;
}