#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <new>
#include "TRW-S.h"

#define private public
#include "typeTruncatedQuadratic2D.h"
#undef private


#define m_D(pix,l)  m_D[(pix)*m_nLabels+(l)]
#define m_V(l1,l2)  m_V[(l1)*m_nLabels+(l2)]


#define MIN(a,b)  (((a) < (b)) ? (a) : (b))
#define MAX(a,b)  (((a) > (b)) ? (a) : (b))
#define TRUNCATE_MIN(a,b) { if ((a) > (b)) (a) = (b); }
#define TRUNCATE_MAX(a,b) { if ((a) < (b)) (a) = (b); }
#define TRUNCATE TRUNCATE_MIN

/////////////////////////////////////////////////////////////////////////////
//                  Operations on vectors (arrays of size K)               //
/////////////////////////////////////////////////////////////////////////////

inline void CopyVector(TRWS::REAL* to, MRF::CostVal* from, int K)
{
    TRWS::REAL* to_finish = to + K;
    do
	{
	    *to ++ = *from ++;
	} while (to < to_finish);
}

inline void AddVector(TRWS::REAL* to, TRWS::REAL* from, int K)
{
    TRWS::REAL* to_finish = to + K;
    do
	{
	    *to ++ += *from ++;
	} while (to < to_finish);
}

inline TRWS::REAL SubtractMin(TRWS::REAL *D, int K)
{
    int k;
    TRWS::REAL delta;
	
    delta = D[0];
    for (k=1; k<K; k++) TRUNCATE(delta, D[k]);
    for (k=0; k<K; k++) D[k] -= delta;

    return delta;
}

// Functions UpdateMessageTYPE (see the paper for details):
//
// - Set Di[ki] := gamma*Di_hat[ki] - M[ki]
// - Set M[kj] := min_{ki} (Di[ki] + V[ki,kj])
// - Normalize message: 
//        delta := min_{kj} M[kj]
//        M[kj] := M[kj] - delta
//        return delta
//
// If dir = 1, then the meaning of i and j is swapped.

///////////////////////////////////////////
//                  L1                   //
///////////////////////////////////////////

inline TRWS::REAL UpdateMessageL1(TRWS::REAL* M, TRWS::REAL* Di_hat, int K, TRWS::REAL gamma, MRF::CostVal lambda, MRF::CostVal smoothMax)
{
    int k;
    TRWS::REAL delta;

    delta = M[0] = gamma*Di_hat[0] - M[0];
    for (k=1; k<K; k++)
	{
	    M[k] = gamma*Di_hat[k] - M[k];
	    TRUNCATE(delta, M[k]);
	    TRUNCATE(M[k], M[k-1] + lambda);
	}

    M[--k] -= delta;
    TRUNCATE(M[k], lambda*smoothMax);
    for (k--; k>=0; k--)
	{
	    M[k] -= delta;
	    TRUNCATE(M[k], M[k+1] + lambda);
	    TRUNCATE(M[k], lambda*smoothMax);
	}

    return delta;
}

////////////////////////////////////////
//               L2                   //
////////////////////////////////////////

inline TRWS::REAL UpdateMessageL2(TRWS::REAL* M, TRWS::REAL* Di_hat, int K, TRWS::REAL gamma, MRF::CostVal lambda, MRF::CostVal smoothMax, void *buf)
{
    TRWS::REAL* Di = (TRWS::REAL*) buf;
    int* parabolas = (int*) ((char*)buf + K*sizeof(TRWS::REAL));
    int* intersections = parabolas + K;
    TypeTruncatedQuadratic2D::Edge* tmp = NULL;

    int k;
    TRWS::REAL delta;

    assert(lambda >= 0);

    Di[0] = gamma*Di_hat[0] - M[0];
    delta = Di[0];
    for (k=1; k<K; k++)
	{
	    Di[k] = gamma*Di_hat[k] - M[k];
	    TRUNCATE(delta, Di[k]);
	}

    if (lambda == 0)
	{
	    for (k=0; k<K; k++) M[k] = 0;
	    return delta;
	}

    tmp->DistanceTransformL2(K, 1, lambda, Di, M, parabolas, intersections);

    for (k=0; k<K; k++)
	{
	    M[k] -= delta;
	    TRUNCATE(M[k], lambda*smoothMax);
	}

    return delta;
}


//////////////////////////////////////////////////
//                FIXED_MATRIX                  //
//////////////////////////////////////////////////

inline TRWS::REAL UpdateMessageFIXED_MATRIX(TRWS::REAL* M, TRWS::REAL* Di_hat, int K, TRWS::REAL gamma, MRF::CostVal lambda, MRF::CostVal* V, void* buf)
{
    TRWS::REAL* Di = (TRWS::REAL*) buf;
    int ki, kj;
    TRWS::REAL delta;

    if (lambda == 0)
	{
	    delta = gamma*Di_hat[0] - M[0];
	    M[0] = 0;
	    for (ki=1; ki<K; ki++) 
		{
		    TRUNCATE(delta, gamma*Di_hat[ki] - M[ki]);
		    M[ki] = 0;
		}
	    return delta;
	}

    for (ki=0; ki<K; ki++)
	{
	    Di[ki] = (gamma*Di_hat[ki] - M[ki]) * (1/(TRWS::REAL)lambda);
	}

    if (lambda > 0)
	{
	    for (kj=0; kj<K; kj++)
		{
		    M[kj] = Di[0] + V[0]; 
		    V ++;
		    for (ki=1; ki<K; ki++)
			{
			    TRUNCATE(M[kj], Di[ki] + V[0]);
			    V ++;
			}
		    M[kj] *= lambda;
		}
	}
    else
	{
	    for (kj=0; kj<K; kj++)
		{
		    M[kj] = Di[0] + V[0]; 
		    V ++;
		    for (ki=1; ki<K; ki++)
			{
			    TRUNCATE_MAX(M[kj], Di[ki] + V[0]);
			    V ++;
			}
		    M[kj] *= lambda;
		}
	}

    delta = M[0];
    for (kj=1; kj<K; kj++) TRUNCATE(delta, M[kj]);
    for (kj=0; kj<K; kj++) M[kj] -= delta;

    return delta;
}

/////////////////////////////////////////////
//                GENERAL                  //
/////////////////////////////////////////////

inline TRWS::REAL UpdateMessageGENERAL(TRWS::REAL* M, TRWS::REAL* Di_hat, int K, TRWS::REAL gamma, int dir, MRF::CostVal* V, void* buf)
{
    TRWS::REAL* Di = (TRWS::REAL*) buf;
    int ki, kj;
    TRWS::REAL delta;

    for (ki=0; ki<K; ki++)
	{
	    Di[ki] = (gamma*Di_hat[ki] - M[ki]);
	}

    if (dir == 0)
	{
	    for (kj=0; kj<K; kj++)
		{
		    M[kj] = Di[0] + V[0]; 
		    V ++;
		    for (ki=1; ki<K; ki++)
			{
			    TRUNCATE(M[kj], Di[ki] + V[0]);
			    V ++;
			}
		}
	}
    else
	{
	    for (kj=0; kj<K; kj++)
		{
		    M[kj] = Di[0] + V[0];
		    V += K;
		    for (ki=1; ki<K; ki++)
			{
			    TRUNCATE(M[kj], Di[ki] + V[0]);
			    V += K;
			}
		    V -= K*K - 1;
		}
	}

    delta = M[0];
    for (kj=1; kj<K; kj++) TRUNCATE(delta, M[kj]);
    for (kj=0; kj<K; kj++) M[kj] -= delta;

    return delta;
}

inline TRWS::REAL UpdateMessageGENERAL(TRWS::REAL* M, TRWS::REAL* Di_hat, int K, TRWS::REAL gamma, TRWS::SmoothCostGeneralFn fn, int i, int j, void* buf)
{
    TRWS::REAL* Di = (TRWS::REAL*) buf;
    int ki, kj;
    TRWS::REAL delta;

    for (ki=0; ki<K; ki++)
	{
	    Di[ki] = (gamma*Di_hat[ki] - M[ki]);
	}

    for (kj=0; kj<K; kj++)
	{
	    M[kj] = Di[0] + fn(i, j, 0, kj); 
	    for (ki=1; ki<K; ki++)
		{
		    delta = Di[ki] + fn(i, j, ki, kj);
		    TRUNCATE(M[kj], delta);
		}
	}

    delta = M[0];
    for (kj=1; kj<K; kj++) TRUNCATE(delta, M[kj]);
    for (kj=0; kj<K; kj++) M[kj] -= delta;

    return delta;
}






TRWS::TRWS(int width, int height, int nLabels,EnergyFunction *eng):MRF(width,height,nLabels,eng)
{
    Allocate();
}
TRWS::TRWS(int nPixels, int nLabels,EnergyFunction *eng):MRF(nPixels,nLabels,eng)
{
    Allocate();
}

TRWS::~TRWS()
{ 
    delete[] m_answer;
    if ( m_needToFreeD ) delete [] m_D;
    if ( m_needToFreeV ) delete [] m_V;
    if ( m_messages ) delete [] m_messages;
    if ( m_DBinary ) delete [] m_DBinary;
    if ( m_horzWeightsBinary ) delete [] m_horzWeightsBinary;
    if ( m_vertWeightsBinary ) delete [] m_vertWeightsBinary;
}


void TRWS::Allocate()
{
    m_type = NONE;
    m_needToFreeV = false;
    m_needToFreeD = false;

    m_D = NULL;
    m_V = NULL;
    m_horzWeights = NULL;
    m_vertWeights = NULL;
    m_horzWeightsBinary = NULL;
    m_vertWeightsBinary = NULL;

    m_DBinary = NULL;
    m_messages = NULL;
    m_messageArraySizeInBytes = 0;

    m_answer = new Label[m_nPixels];
}

void TRWS::clearAnswer()
{
    memset(m_answer, 0, m_nPixels*sizeof(Label));
    if (m_messages)
	{
	    memset(m_messages, 0, m_messageArraySizeInBytes);
	}
}


MRF::EnergyVal TRWS::smoothnessEnergy()
{
    EnergyVal eng = (EnergyVal) 0;
    EnergyVal weight;
    int x,y,pix;

    if ( m_grid_graph )
	{
	    if ( m_smoothType != FUNCTION  )
		{
		    for ( y = 0; y < m_height; y++ )
			for ( x = 1; x < m_width; x++ )
			    {
				pix    = x+y*m_width;
				weight = m_varWeights ? m_horzWeights[pix-1] :  1;
				eng = eng + m_V(m_answer[pix],m_answer[pix-1])*weight;
			    }

		    for ( y = 1; y < m_height; y++ )
			for ( x = 0; x < m_width; x++ )
			    {
				pix = x+y*m_width;
				weight = m_varWeights ? m_vertWeights[pix-m_width] :  1;
				eng = eng + m_V(m_answer[pix],m_answer[pix-m_width])*weight;
			    }
		}
	    else
		{
		    for ( y = 0; y < m_height; y++ )
			for ( x = 1; x < m_width; x++ )
			    {
				pix = x+y*m_width;
				eng = eng + m_smoothFn(pix,pix-1,m_answer[pix],m_answer[pix-1]);
			    }

		    for ( y = 1; y < m_height; y++ )
			for ( x = 0; x < m_width; x++ )
			    {
				pix = x+y*m_width;
				eng = eng + m_smoothFn(pix,pix-m_width,m_answer[pix],m_answer[pix-m_width]);
			    }
		}
	}
    else
	{
	    // not implemented
	}

    return(eng);
}



MRF::EnergyVal TRWS::dataEnergy()
{
    EnergyVal eng = (EnergyVal) 0;

    
    if ( m_dataType == ARRAY) 
	{
	    for ( int i = 0; i < m_nPixels; i++ )
		eng = eng + m_D(i,m_answer[i]);
	}
    else
	{
	    for ( int i = 0; i < m_nPixels; i++ )
		eng = eng + m_dataFn(i,m_answer[i]);
	}
    return(eng);

}


void TRWS::setData(DataCostFn dcost)
{
    int i, k;

    m_dataFn = dcost;
    CostVal* ptr;
    m_D = new CostVal[m_nPixels*m_nLabels];

    for (ptr=m_D, i=0; i<m_nPixels; i++)
	for (k=0; k<m_nLabels; k++, ptr++)
	    {
		*ptr = m_dataFn(i,k);
	    }
    m_needToFreeD = true;
}

void TRWS::setData(CostVal* data)
{
    m_D = data;
    m_needToFreeD = false;
}


void TRWS::setSmoothness(SmoothCostGeneralFn cost)
{
    assert(m_horzWeights == NULL && m_vertWeights == NULL && m_V == NULL);

    int x, y, i, ki, kj;
    CostVal* ptr;

    m_smoothFn = cost;
    m_type = GENERAL;

    if (!m_allocateArrayForSmoothnessCostFn) return;

    // try to cache all the function values in an array for efficiency
    m_V = new(std::nothrow) CostVal[2*m_nPixels*m_nLabels*m_nLabels];
    if (!m_V) return; // if not enough space, just call the function directly

    m_needToFreeV = true;

    for (ptr=m_V,i=0,y=0; y<m_height; y++)
	for (x=0; x<m_width; x++, i++)
	    {
		if (x < m_width-1)
		    {
			for (kj=0; kj<m_nLabels; kj++)
			    for (ki=0; ki<m_nLabels; ki++)
				{
				    *ptr++ = cost(i,i+1,ki,kj);
				}
		    }
		else ptr += m_nLabels*m_nLabels;

		if (y < m_height-1)
		    {
			for (kj=0; kj<m_nLabels; kj++)
			    for (ki=0; ki<m_nLabels; ki++)
				{
				    *ptr++ = cost(i,i+m_width,ki,kj);
				}
		    }
		else ptr += m_nLabels*m_nLabels;
	    }
}
void TRWS::setSmoothness(CostVal* V)
{
    m_type = FIXED_MATRIX;
    m_V = V;
}


void TRWS::setSmoothness(int smoothExp,CostVal smoothMax, CostVal lambda)
{
    assert(smoothExp == 1 || smoothExp == 2);
    assert(lambda >= 0);

    m_type = (smoothExp == 1) ? L1 : L2;

    int ki, kj;
    CostVal cost;

    m_needToFreeV = true;

    m_V = new CostVal[m_nLabels*m_nLabels];

    for (ki=0; ki<m_nLabels; ki++)
	for (kj=ki; kj<m_nLabels; kj++)
	    {
		cost = (CostVal) ((smoothExp == 1) ? kj - ki : (kj - ki)*(kj - ki));
		if (cost > smoothMax) cost = smoothMax;
		m_V[ki*m_nLabels + kj] = m_V[kj*m_nLabels + ki] = cost*lambda;
	    }

    m_smoothMax = smoothMax;
    m_lambda = lambda;
}


void TRWS::setCues(CostVal* hCue, CostVal* vCue)
{
    m_horzWeights = hCue;
    m_vertWeights  = vCue;
}


void TRWS::initializeAlg()
{
    assert(m_type != NONE);

    int i;

    // determine type
    if (m_type == L1 && m_nLabels == 2)
	{
	    m_type = BINARY;
	}

    // allocate messages
    int messageNum = (m_type == BINARY) ? 4*m_nPixels : 4*m_nPixels*m_nLabels;
    m_messageArraySizeInBytes = messageNum*sizeof(REAL);
    m_messages = new REAL[messageNum];
    memset(m_messages, 0, messageNum*sizeof(REAL));

    if (m_type == BINARY)
	{
	    assert(m_DBinary == NULL && m_horzWeightsBinary == NULL && m_horzWeightsBinary == NULL);
	    m_DBinary = new CostVal[m_nPixels];
	    m_horzWeightsBinary = new CostVal[m_nPixels];
	    m_vertWeightsBinary = new CostVal[m_nPixels];

	    if ( m_dataType == ARRAY)
		{
		    for (i=0; i<m_nPixels; i++)
			{
			    m_DBinary[i] = m_D[2*i+1] - m_D[2*i];
			}
		}
	    else
		{
		    for (i=0; i<m_nPixels; i++)
			{
			    m_DBinary[i] = m_dataFn(i,1) - m_dataFn(i,0);
			}
		}

	    assert(m_V[0] == 0 && m_V[1] == m_V[2] && m_V[3] == 0);
	    for (i=0; i<m_nPixels; i++)
		{
		    m_horzWeightsBinary[i] = (m_varWeights) ? m_V[1]*m_horzWeights[i] : m_V[1];
		    m_vertWeightsBinary[i] = (m_varWeights) ? m_V[1]*m_vertWeights[i] : m_V[1];
		}
	}
}

void TRWS::optimizeAlg(int nIterations)
{
    assert(m_type != NONE);

    if (m_grid_graph)
	{
	    switch (m_type)
		{
		case L1:            optimize_GRID_L1(nIterations);           break;
		case L2:            optimize_GRID_L2(nIterations);           break;
		case FIXED_MATRIX:  optimize_GRID_FIXED_MATRIX(nIterations); break;
		case GENERAL:       optimize_GRID_GENERAL(nIterations);      break;
		case BINARY:        optimize_GRID_BINARY(nIterations);       break;
		default: assert(0); exit(1);
		}
	}
    else {printf("\nNot implemented for general graphs yet, exiting!");exit(1);}

    //	printf("lower bound = %f\n", m_lowerBound);

    ////////////////////////////////////////////////
    //          computing solution                //
    ////////////////////////////////////////////////

    if (m_type != BINARY)
	{
	    int x, y, n, K = m_nLabels;
	    CostVal* D_ptr;
	    REAL* M_ptr;
	    REAL* Di;
	    REAL delta;
	    int ki, kj;

	    Di = new REAL[K];

	    n = 0;
	    D_ptr = m_D;
	    M_ptr = m_messages;

	    for (y=0; y<m_height; y++)
		for (x=0; x<m_width; x++, D_ptr+=K, M_ptr+=2*K, n++)
		    {
			CopyVector(Di, D_ptr, K);

			if (m_type == GENERAL)
			    {
				if (m_V)
				    {
					CostVal* ptr = m_V + 2*(x+y*m_width-1)*K*K;
					if (x > 0)
					    {
						kj = m_answer[n-1];
						for (ki=0; ki<K; ki++)
						    {
							Di[ki] += ptr[kj + ki*K];
						    }
					    }
					ptr -= (2*m_width-3)*K*K;
					if (y > 0)
					    {
						kj = m_answer[n-m_width];
						for (ki=0; ki<K; ki++)
						    {
							Di[ki] += ptr[kj + ki*K];
						    }
					    }
				    }
				else
				    {
					if (x > 0)
					    {
						kj = m_answer[n-1];
						for (ki=0; ki<K; ki++)
						    {
							Di[ki] += m_smoothFn(n, n-1, ki, kj);
						    }
					    }
					if (y > 0)
					    {
						kj = m_answer[n-m_width];
						for (ki=0; ki<K; ki++)
						    {
							Di[ki] += m_smoothFn(n, n-m_width, ki, kj);
						    }
					    }
				    }
			    }
			else // m_type == L1, L2 or FIXED_MATRIX
			    {
				if (x > 0)
				    {
					kj = m_answer[n-1];
					CostVal lambda = (m_varWeights) ? m_horzWeights[n-1] : 1;
					for (ki=0; ki<K; ki++)
					    {
						Di[ki] += lambda*m_V[kj*K + ki];
					    }
				    }
				if (y > 0)
				    {
					kj = m_answer[n-m_width];
					CostVal lambda = (m_varWeights) ? m_vertWeights[n-m_width] : 1;
					for (ki=0; ki<K; ki++)
					    {
						Di[ki] += lambda*m_V[kj*K + ki];
					    }
				    }
			    }

			if (x < m_width-1) AddVector(Di, M_ptr, K); // message (x+1,y)->(x,y)
			if (y < m_height-1) AddVector(Di, M_ptr+K, K); // message (x,y+1)->(x,y)

			// compute min
			delta = Di[0];
			m_answer[n] = 0;
			for (ki=1; ki<K; ki++)
			    {
				if (delta > Di[ki])
				    {
					delta = Di[ki];
					m_answer[n] = ki;
				    }
			    }
		    }

	    delete [] Di;
	}
    else // m_type == BINARY
	{
	    int x, y, n;
	    REAL* M_ptr;
	    REAL Di;

	    n = 0;
	    M_ptr = m_messages;

	    for (y=0; y<m_height; y++)
		for (x=0; x<m_width; x++, M_ptr+=2, n++)
		    {
			Di = m_DBinary[n];
			if (x > 0) Di += (m_answer[n-1] == 0)       ? m_horzWeightsBinary[n-1]       : -m_horzWeightsBinary[n-1];
			if (y > 0) Di += (m_answer[n-m_width] == 0) ? m_vertWeightsBinary[n-m_width] : -m_vertWeightsBinary[n-m_width];

			if (x < m_width-1)  Di += M_ptr[0]; // message (x+1,y)->(x,y)
			if (y < m_height-1) Di += M_ptr[1]; // message (x,y+1)->(x,y)

			// compute min
			m_answer[n] = (Di >= 0) ? 0 : 1;
		    }
	}
}

void TRWS::optimize_GRID_L1(int nIterations)
{
    int x, y, n, K = m_nLabels;
    CostVal* D_ptr;
    REAL* M_ptr;
    REAL* Di;

    Di = new REAL[K];

    for ( ; nIterations > 0; nIterations --)
	{
	    // forward pass
	    n = 0;
	    D_ptr = m_D;
	    M_ptr = m_messages;

	    for (y=0; y<m_height; y++)
		for (x=0; x<m_width; x++, D_ptr+=K, M_ptr+=2*K, n++)
		    {
			CopyVector(Di, D_ptr, K);
			if (x > 0) AddVector(Di, M_ptr-2*K, K); // message (x-1,y)->(x,y)
			if (y > 0) AddVector(Di, M_ptr-(2*m_width-1)*K, K); // message (x,y-1)->(x,y)
			if (x < m_width-1) AddVector(Di, M_ptr, K); // message (x+1,y)->(x,y)
			if (y < m_height-1) AddVector(Di, M_ptr+K, K); // message (x,y+1)->(x,y)

			if (x < m_width-1) 
			    {
				CostVal lambda = (m_varWeights) ? m_lambda*m_horzWeights[n] : m_lambda;
				UpdateMessageL1(M_ptr, Di, K, 0.5, lambda, m_smoothMax);
			    }
			if (y < m_height-1) 
			    {
				CostVal lambda = (m_varWeights) ? m_lambda*m_vertWeights[n] : m_lambda;
				UpdateMessageL1(M_ptr+K, Di, K, 0.5, lambda, m_smoothMax);
			    }
		    }

	    // backward pass
	    m_lowerBound = 0;

	    n --;
	    D_ptr -= K;
	    M_ptr -= 2*K;

	    for (y=m_height-1; y>=0; y--)
		for (x=m_width-1; x>=0; x--, D_ptr-=K, M_ptr-=2*K, n--)
		    {
			CopyVector(Di, D_ptr, K);
			if (x > 0) AddVector(Di, M_ptr-2*K, K); // message (x-1,y)->(x,y)
			if (y > 0) AddVector(Di, M_ptr-(2*m_width-1)*K, K); // message (x,y-1)->(x,y)
			if (x < m_width-1) AddVector(Di, M_ptr, K); // message (x+1,y)->(x,y)
			if (y < m_height-1) AddVector(Di, M_ptr+K, K); // message (x,y+1)->(x,y)

			m_lowerBound += SubtractMin(Di, K);

			if (x > 0) 
			    {
				CostVal lambda = (m_varWeights) ? m_lambda*m_horzWeights[n-1] : m_lambda;
				m_lowerBound += UpdateMessageL1(M_ptr-2*K, Di, K, 0.5, lambda, m_smoothMax);
			    }
			if (y > 0) 
			    {
				CostVal lambda = (m_varWeights) ? m_lambda*m_vertWeights[n-m_width] : m_lambda;
				m_lowerBound += UpdateMessageL1(M_ptr-(2*m_width-1)*K, Di, K, 0.5, lambda, m_smoothMax);
			    }
		    }
	}

    delete [] Di;
}

void TRWS::optimize_GRID_L2(int nIterations)
{
    int x, y, n, K = m_nLabels;
    CostVal* D_ptr;
    REAL* M_ptr;
    REAL* Di;
    void* buf;

    Di = new REAL[K];
    buf = new char[(2*K+1)*sizeof(int) + K*sizeof(REAL)];

    for ( ; nIterations > 0; nIterations --)
	{
	    // forward pass
	    n = 0;
	    D_ptr = m_D;
	    M_ptr = m_messages;

	    for (y=0; y<m_height; y++)
		for (x=0; x<m_width; x++, D_ptr+=K, M_ptr+=2*K, n++)
		    {
			CopyVector(Di, D_ptr, K);
			if (x > 0) AddVector(Di, M_ptr-2*K, K); // message (x-1,y)->(x,y)
			if (y > 0) AddVector(Di, M_ptr-(2*m_width-1)*K, K); // message (x,y-1)->(x,y)
			if (x < m_width-1) AddVector(Di, M_ptr, K); // message (x+1,y)->(x,y)
			if (y < m_height-1) AddVector(Di, M_ptr+K, K); // message (x,y+1)->(x,y)

			if (x < m_width-1) 
			    {
				CostVal lambda = (m_varWeights) ? m_lambda*m_horzWeights[n] : m_lambda;
				UpdateMessageL2(M_ptr, Di, K, 0.5, lambda, m_smoothMax, buf);
			    }
			if (y < m_height-1) 
			    {
				CostVal lambda = (m_varWeights) ? m_lambda*m_vertWeights[n] : m_lambda;
				UpdateMessageL2(M_ptr+K, Di, K, 0.5, lambda, m_smoothMax, buf);
			    }
		    }

	    // backward pass
	    m_lowerBound = 0;

	    n --;
	    D_ptr -= K;
	    M_ptr -= 2*K;

	    for (y=m_height-1; y>=0; y--)
		for (x=m_width-1; x>=0; x--, D_ptr-=K, M_ptr-=2*K, n--)
		    {
			CopyVector(Di, D_ptr, K);
			if (x > 0) AddVector(Di, M_ptr-2*K, K); // message (x-1,y)->(x,y)
			if (y > 0) AddVector(Di, M_ptr-(2*m_width-1)*K, K); // message (x,y-1)->(x,y)
			if (x < m_width-1) AddVector(Di, M_ptr, K); // message (x+1,y)->(x,y)
			if (y < m_height-1) AddVector(Di, M_ptr+K, K); // message (x,y+1)->(x,y)

			m_lowerBound += SubtractMin(Di, K);

			if (x > 0) 
			    {
				CostVal lambda = (m_varWeights) ? m_lambda*m_horzWeights[n-1] : m_lambda;
				m_lowerBound += UpdateMessageL2(M_ptr-2*K, Di, K, 0.5, lambda, m_smoothMax, buf);
			    }
			if (y > 0) 
			    {
				CostVal lambda = (m_varWeights) ? m_lambda*m_vertWeights[n-m_width] : m_lambda;
				m_lowerBound += UpdateMessageL2(M_ptr-(2*m_width-1)*K, Di, K, 0.5, lambda, m_smoothMax, buf);
			    }
		    }
	}

    delete [] Di;
    delete [] (char *)buf;
}


void TRWS::optimize_GRID_BINARY(int nIterations)
{
    int x, y, n;
    REAL* M_ptr;
    REAL Di;

    for ( ; nIterations > 0; nIterations --)
	{
	    // forward pass
	    n = 0;
	    M_ptr = m_messages;

	    for (y=0; y<m_height; y++)
		for (x=0; x<m_width; x++, M_ptr+=2, n++)
		    {
			Di = m_DBinary[n];
			if (x > 0) Di += M_ptr[-2]; // message (x-1,y)->(x,y)
			if (y > 0) Di += M_ptr[-2*m_width+1]; // message (x,y-1)->(x,y)
			if (x < m_width-1) Di += M_ptr[0]; // message (x+1,y)->(x,y)
			if (y < m_height-1) Di += M_ptr[1]; // message (x,y+1)->(x,y)

			REAL DiScaled = Di * 0.5;
			if (x < m_width-1) 
			    {
				Di = DiScaled - M_ptr[0];
				CostVal lambda = m_horzWeightsBinary[n];
				if (lambda < 0) { Di = -Di; lambda = -lambda; }
				if (Di > lambda) M_ptr[0] = lambda;
				else             M_ptr[0] = (Di < -lambda) ? -lambda : Di;
			    }
			if (y < m_height-1) 
			    {
				Di = DiScaled - M_ptr[1];
				CostVal lambda = m_vertWeightsBinary[n];
				if (lambda < 0) { Di = -Di; lambda = -lambda; }
				if (Di > lambda) M_ptr[1] = lambda;
				else             M_ptr[1] = (Di < -lambda) ? -lambda : Di;
			    }
		    }

	    // backward pass
	    n --;
	    M_ptr -= 2;

	    for (y=m_height-1; y>=0; y--)
		for (x=m_width-1; x>=0; x--, M_ptr-=2, n--)
		    {
			Di = m_DBinary[n];
			if (x > 0) Di += M_ptr[-2]; // message (x-1,y)->(x,y)
			if (y > 0) Di += M_ptr[-2*m_width+1]; // message (x,y-1)->(x,y)
			if (x < m_width-1) Di += M_ptr[0]; // message (x+1,y)->(x,y)
			if (y < m_height-1) Di += M_ptr[1]; // message (x,y+1)->(x,y)

			REAL DiScaled = Di * 0.5;
			if (x > 0) 
			    {
				Di = DiScaled - M_ptr[-2];
				CostVal lambda = m_horzWeightsBinary[n-1];
				if (lambda < 0) { Di = -Di; lambda = -lambda; }
				if (Di > lambda) M_ptr[-2] = lambda;
				else             M_ptr[-2] = (Di < -lambda) ? -lambda : Di;
			    }
			if (y > 0) 
			    {
				Di = DiScaled - M_ptr[-2*m_width+1];
				CostVal lambda = m_vertWeightsBinary[n-m_width];
				if (lambda < 0) { Di = -Di; lambda = -lambda; }
				if (Di > lambda) M_ptr[-2*m_width+1] = lambda;
				else             M_ptr[-2*m_width+1] = (Di < -lambda) ? -lambda : Di;
			    }
		    }
	}

    m_lowerBound = 0;
}

void TRWS::optimize_GRID_FIXED_MATRIX(int nIterations)
{
    int x, y, n, K = m_nLabels;
    CostVal* D_ptr;
    REAL* M_ptr;
    REAL* Di;
    void* buf;

    Di = new REAL[K];
    buf = new REAL[K];

    for ( ; nIterations > 0; nIterations --)
	{
	    // forward pass
	    n = 0;
	    D_ptr = m_D;
	    M_ptr = m_messages;

	    for (y=0; y<m_height; y++)
		for (x=0; x<m_width; x++, D_ptr+=K, M_ptr+=2*K, n++)
		    {
			CopyVector(Di, D_ptr, K);
			if (x > 0) AddVector(Di, M_ptr-2*K, K); // message (x-1,y)->(x,y)
			if (y > 0) AddVector(Di, M_ptr-(2*m_width-1)*K, K); // message (x,y-1)->(x,y)
			if (x < m_width-1) AddVector(Di, M_ptr, K); // message (x+1,y)->(x,y)
			if (y < m_height-1) AddVector(Di, M_ptr+K, K); // message (x,y+1)->(x,y)

			if (x < m_width-1) 
			    {
				CostVal lambda = (m_varWeights) ? m_horzWeights[n] : 1;
				UpdateMessageFIXED_MATRIX(M_ptr, Di, K, 0.5, lambda, m_V, buf);
			    }
			if (y < m_height-1) 
			    {
				CostVal lambda = (m_varWeights) ? m_vertWeights[n] : 1;
				UpdateMessageFIXED_MATRIX(M_ptr+K, Di, K, 0.5, lambda, m_V, buf);
			    }
		    }

	    // backward pass
	    m_lowerBound = 0;

	    n --;
	    D_ptr -= K;
	    M_ptr -= 2*K;

	    for (y=m_height-1; y>=0; y--)
		for (x=m_width-1; x>=0; x--, D_ptr-=K, M_ptr-=2*K, n--)
		    {
			CopyVector(Di, D_ptr, K);
			if (x > 0) AddVector(Di, M_ptr-2*K, K); // message (x-1,y)->(x,y)
			if (y > 0) AddVector(Di, M_ptr-(2*m_width-1)*K, K); // message (x,y-1)->(x,y)
			if (x < m_width-1) AddVector(Di, M_ptr, K); // message (x+1,y)->(x,y)
			if (y < m_height-1) AddVector(Di, M_ptr+K, K); // message (x,y+1)->(x,y)

			m_lowerBound += SubtractMin(Di, K);

			if (x > 0) 
			    {
				CostVal lambda = (m_varWeights) ? m_horzWeights[n-1] : 1;
				m_lowerBound += UpdateMessageFIXED_MATRIX(M_ptr-2*K, Di, K, 0.5, lambda, m_V, buf);
			    }
			if (y > 0) 
			    {
				CostVal lambda = (m_varWeights) ? m_vertWeights[n-m_width] : 1;
				m_lowerBound += UpdateMessageFIXED_MATRIX(M_ptr-(2*m_width-1)*K, Di, K, 0.5, lambda, m_V, buf);
			    }
		    }
	}

    delete [] Di;
    delete [] (REAL *)buf;
}

void TRWS::optimize_GRID_GENERAL(int nIterations)
{
    int x, y, n, K = m_nLabels;
    CostVal* D_ptr;
    REAL* M_ptr;
    REAL* Di;
    void* buf;

    Di = new REAL[K];
    buf = new REAL[K];

    for ( ; nIterations > 0; nIterations --)
	{
	    // forward pass
	    n = 0;
	    D_ptr = m_D;
	    M_ptr = m_messages;
	    CostVal* V_ptr = m_V;

	    for (y=0; y<m_height; y++)
		for (x=0; x<m_width; x++, D_ptr+=K, M_ptr+=2*K, V_ptr+=2*K*K, n++)
		    {
			CopyVector(Di, D_ptr, K);
			if (x > 0) AddVector(Di, M_ptr-2*K, K); // message (x-1,y)->(x,y)
			if (y > 0) AddVector(Di, M_ptr-(2*m_width-1)*K, K); // message (x,y-1)->(x,y)
			if (x < m_width-1) AddVector(Di, M_ptr, K); // message (x+1,y)->(x,y)
			if (y < m_height-1) AddVector(Di, M_ptr+K, K); // message (x,y+1)->(x,y)

			if (x < m_width-1) 
			    {
				if (m_V) UpdateMessageGENERAL(M_ptr, Di, K, 0.5, /* forward dir*/ 0, V_ptr, buf);
				else     UpdateMessageGENERAL(M_ptr, Di, K, 0.5,   m_smoothFn, n, n+1,      buf);
			    }
			if (y < m_height-1) 
			    {
				if (m_V) UpdateMessageGENERAL(M_ptr+K, Di, K, 0.5, /* forward dir*/ 0, V_ptr+K*K, buf);
				else     UpdateMessageGENERAL(M_ptr+K, Di, K, 0.5,   m_smoothFn, n, n+m_width,    buf);
			    }
		    }

	    // backward pass
	    m_lowerBound = 0;

	    n --;
	    D_ptr -= K;
	    M_ptr -= 2*K;
	    V_ptr -= 2*K*K;

	    for (y=m_height-1; y>=0; y--)
		for (x=m_width-1; x>=0; x--, D_ptr-=K, M_ptr-=2*K, V_ptr-=2*K*K, n--)
		    {
			CopyVector(Di, D_ptr, K);
			if (x > 0) AddVector(Di, M_ptr-2*K, K); // message (x-1,y)->(x,y)
			if (y > 0) AddVector(Di, M_ptr-(2*m_width-1)*K, K); // message (x,y-1)->(x,y)
			if (x < m_width-1) AddVector(Di, M_ptr, K); // message (x+1,y)->(x,y)
			if (y < m_height-1) AddVector(Di, M_ptr+K, K); // message (x,y+1)->(x,y)

			// normalize Di, update lower bound
			m_lowerBound += SubtractMin(Di, K);

			if (x > 0) 
			    {
				if (m_V) m_lowerBound += UpdateMessageGENERAL(M_ptr-2*K, Di, K, 0.5, /* backward dir */ 1, V_ptr-2*K*K, buf);
				else     m_lowerBound += UpdateMessageGENERAL(M_ptr-2*K, Di, K, 0.5,   m_smoothFn, n, n-1,              buf);
			    }
			if (y > 0)
			    {
				if (m_V) m_lowerBound += UpdateMessageGENERAL(M_ptr-(2*m_width-1)*K, Di, K, 0.5, /* backward dir */ 1, V_ptr-(2*m_width-1)*K*K, buf);
				else     m_lowerBound += UpdateMessageGENERAL(M_ptr-(2*m_width-1)*K, Di, K, 0.5,   m_smoothFn, n, n-m_width,                    buf);
			    }
		    }
	}

    delete [] Di;
    delete [] (REAL *)buf;
}