#include "DynamicProgramSearch.h"
#include <list>

using std::list;

void dynamicProgramY(Mat& disMat, int n /*= 2*/)
{
    for (int y = 1; y < disMat.rows; ++y)
    {
        for (int x = 0; x < disMat.cols; ++x)
        {
            // find max neighbor
            int sy = y - 1;
            float maxVal = -1;
            int sj = x - n >= 0 ? x - n : 0;
            int ej = x + n <= disMat.cols - 1 ? x + n : disMat.cols - 1;
            for (int j = sj; j <= ej; j++)
            {
                float val = disMat.at<float>(sy, j);
                if (maxVal < val)
                {
                    maxVal = val;
                }
            }
            
            disMat.at<float>(y, x) += maxVal;
        }
    }
}

void dynamicProgramX(Mat& disMat, int n /*= 2*/)
{
    for (int x = 1; x < disMat.cols; ++x)
    {
        int sx = x - 1;
        for (int y = 0; y < disMat.rows; ++y)
        {
            // find max neighbor
            float maxVal = -1;
            int si = max(y - n, 0);
            int ei = min(y + n, disMat.rows - 1);
            for (int i = si; i <= ei; i++)
            {
                float val = disMat.at<float>(i, sx);
                if (maxVal < val)
                {
                    maxVal = val;
                }
            }

            disMat.at<float>(y, x) += maxVal;
        }
    }
}

void findMaxXPath(const Mat& disMat, vector<int>& vec, vector<float>& energyVec,
    int n /*= 2*/)
{
    int sy = 0;
    int ey = disMat.rows - 1;
    vec.resize(disMat.cols, -1);
    energyVec.resize(disMat.cols, -1);
    for (int x = disMat.cols - 1; x >= 0; --x)
    {
        float maxVal = -1;
        int maxIdx = -1;
        for (int y = sy; y <= ey; ++y)
        {
            float val = disMat.at<float>(y, x);
            if (maxVal < val)
            {
                maxVal = val;
                maxIdx = y;
            }
        }
        vec[x] = maxIdx;
        energyVec[x] = maxVal;
        sy = maxIdx - n;
        ey = maxIdx + n;
        sy = (sy >= 0 ? sy : 0);
        ey = (ey < disMat.rows ? ey : disMat.rows - 1);
    }
}

void findMaxYPath(const Mat& disMat, vector<int>& vec, vector<float>& energyVec, int n /*= 2*/)
{
    int sx = 0;
    int ex = disMat.cols - 1;
    vec.resize(disMat.rows, -1);
    energyVec.resize(disMat.rows, -1);
    for (int y = disMat.rows - 1; y >= 0; --y)
    {
        float maxVal = -1;
        int maxIdx = -1;
        for (int x = sx; x <= ex; ++x)
        {
            float val = disMat.at<float>(y, x);
            if (maxVal < val)
            {
                maxVal = val;
                maxIdx = x;
            }
        }
        vec[y] = maxIdx;
        energyVec[y] = maxVal;
        sx = maxIdx - n;
        ex = maxIdx + n;
        sx = (sx >= 0 ? sx : 0);
        ex = (ex < disMat.cols ? ex : disMat.cols - 1);
    }
}

void recoverPathEnergy(vector<float>& vec)
{
    assert(!vec.empty());

    auto i = vec.rbegin();
    auto j = i;
    ++j;
    while (j != vec.rend())
    {
        *i = *i - *j;
        ++i;
        ++j;
    }
}

void dynamicProgramYWithSmooth(Mat& disMat,
    int n /*= 2*/, float sw /*= 1.0*/,
    const Mat& smoothMat /*= Mat()*/)
{
    for (int y = 1; y < disMat.rows; ++y)
    {
        for (int x = 0; x < disMat.cols; ++x)
        {
            // find max neighbor
            int sy = y - 1;
            float maxVal = -1;
            int sj = x - n >= 0 ? x - n : 0;
            int ej = x + n <= disMat.cols - 1 ? x + n : disMat.cols - 1;
            float baseSVal = 0;
            if (!smoothMat.empty())
            {
                baseSVal = smoothMat.at<float>(y, x);
            }
            for (int j = sj; j <= ej; j++)
            {
                float disVal = disMat.at<float>(sy, j);
                float sVal = 0;
                if (!smoothMat.empty())
                {
                    sVal = smoothMat.at<float>(sy, j);
                }
                float energyVal = disVal + (255 - abs(sVal - baseSVal))*sw;
                if (maxVal < energyVal)
                {
                    maxVal = energyVal;
                }
            }

            disMat.at<float>(y, x) += maxVal;
        }
    }
}

void dynamicProgramYWithSmooth(Mat& disMat0, Mat& disMat1,
	const Mat& smoothMat0,
	const Mat& smoothMat1,
	int targetWidth,
	int n /* = 2 */,
	float sw /* = 1.0 */)
{
	assert(targetWidth <= disMat0.cols + disMat1.cols);

	for (int y = 1; y < disMat0.rows; ++y)
	{
		int sx = 0;
		int ex = disMat0.cols;
		if (targetWidth < disMat0.cols)
		{
			sx = disMat0.cols - targetWidth;
		}
		if (targetWidth > disMat1.cols)
		{
			ex = disMat0.cols - (targetWidth - disMat1.cols);
		}

		for (int x = sx; x < ex; ++x)
		{
			// find max neighbor
			int sy = y - 1;
			float maxVal = -1;
			int sj0 = x - n >= 0 ? x - n : 0;
			int ej0 = x + n <= disMat0.cols - 1 ? x + n : disMat0.cols - 1;
			int sj1 = sj0 + targetWidth - disMat0.cols;
			int ej1 = ej0 + targetWidth - disMat0.cols;
			if (sj1 < 0)
			{
				sj0 += -sj1;
			}
			if (ej1 > disMat1.cols - 1)
			{
				ej0 -= ej1 - disMat1.cols + 1;
			}
			float baseSVal0 = 0;
			if (!smoothMat0.empty())
			{
				baseSVal0 = smoothMat0.at<float>(y, x);
			}
			float baseSVal1 = 0;
			if (!smoothMat1.empty())
			{
				baseSVal1 = smoothMat1.at<float>(y, x + targetWidth - smoothMat0.cols);
			}
			for (int j = sj0; j <= ej0; j++)
			{
				float disVal0 = disMat0.at<float>(sy, j);
				int j1 = j + targetWidth - disMat0.cols;
				float disVal1 = disMat1.at<float>(sy, j1);
				float sVal0 = 0;
				if (!smoothMat0.empty())
				{
					sVal0 = smoothMat0.at<float>(sy, j);
				}
				float sVal1 = 0;
				if (!smoothMat1.empty())
				{
					sVal1 = smoothMat1.at<float>(sy, j1);
				}
				float energyVal = disVal0 + disVal1 +
					(255 - abs(sVal0 - baseSVal0))*sw +
					(255 - abs(sVal1 - baseSVal1))*sw;
				if (maxVal < energyVal)
				{
					maxVal = energyVal;
				}
			}

			disMat0.at<float>(y, x) += maxVal;
		}
	}
}