#include <string>
#include "ximage.h"

#include "simpleimage.h"
#include "simpleimage4cximage.h"
#include "simpleimageoperation.h"

#include "taucsaddon.h"

#include <limits>
#include <map>
#include <iostream>

#if !defined(SAFE_DELETE)
#define SAFE_DELETE(pointer, operation){if (pointer != NULL) {operation (##pointer##); pointer = NULL;}}
#endif

////////////////////////////////////////////////////////////////////////////////
void YUVtoRGB( CSimpleImagef &dest)
{
	for (int y = 0; y < dest._height; ++ y)
	{
		for (int x = 0; x < dest._width; ++ x)
		{
			float *p = dest.at(x, y);

			float Y = p[0];
			float U = p[1];
			float V = p[2];

			float R = ( Y + 1.403f * V);
			float G = ( Y - 0.344f * U - 0.714f * V);
			float B = ( Y + 1.770f * U);

			p[0] = min(255,max(0,B));
			p[1] = min(255,max(0,G));
			p[2] = min(255,max(0,R));			
		}
	}
}

////////////////////////////////////////////////////////////////////////////////
void RGBtoYUV(CSimpleImagef &dest)
{
	for (int y = 0; y < dest._height; ++ y)
	{
		for (int x = 0; x < dest._width; ++ x)
		{
			float *p = dest.at(x, y);
			
			float B = p[0];
			float G = p[1];
			float R = p[2];

			float Y = p[0] = (0.299f * R + 0.587f * G + 0.114f * B);
			p[1] = ((B-Y) * 0.565f);
			p[2] = ((R-Y) * 0.713f);
		}
	}
}

void divergence(CSimpleImagef &src, CSimpleImagef &div, int channel)
{
	CSimpleImagef grad;
	grad.set_size(src._width, src._height, 2);

	for (int y = 0; y < src._height - 1; y++)
	{
		for (int x = 0; x < src._width - 1; x++)
		{
			grad.at(x, y)[0] = src.at(x+1,y)[channel] - src.at(x,y)[channel];
			grad.at(x, y)[1] = src.at(x,y+1)[channel] - src.at(x,y)[channel];
		}
	}

	div.set_size(src._width, src._height, 1);
	
	for (int y = 1; y < src._height - 1; ++ y)
	{
		for (int x = 1; x < src._width - 1; ++ x)
		{
			div.at(x, y)[0] = grad.at(x, y)[0] - grad.at(x - 1, y)[0] + 
				grad.at(x, y)[1] - grad.at(x, y - 1)[1];
		}
	}
}

bool build_discrete_poisson(CSimpleImagef &src, CSimpleImagef &dest, CSimpleImageb &mask,
							taucs_ccs_matrix **A, double **b)
{
	//////////////////////////////////////////////////////////////////////////
	int variablecount = 0;

	// in the masked area, each pixel corresponds to a variable, we assign a unique index to each variable.
	std::map<int, int> pixelvariable;	// variable indexer

	int height = src._height;
	int width = src._width;

	for (int y = 1; y < height - 1; ++ y)
	{
		for (int x = 1; x < width - 1; ++ x)
		{
			if (mask.at(x, y)[0] > 0)
			{
				int pixelindex = y * width + x;
				pixelvariable[pixelindex] = variablecount;
				++ variablecount;
			}
		}
	}

	if (variablecount == 0)
	{
		std::cout << "error: No masked pixels found (mask color is non-black)\n";
		return false;
	}

	std::cout << "Solver::solve: Solving " << width << "x" << height << " with " << variablecount << " unknowns" << std::endl;
	//////////////////////////////////////////////////////////////////////////
	// taucs is spare linear system solver.

	// claim matrix memory.
	taucs_ccs_matrix* pA = *A;
	double* pb = *b;
	SAFE_DELETE(pA, taucs_ccs_free)
	SAFE_DELETE(pb, delete[])

	// 4 neighbors + pixelself, most 5 non-zero elements for each pixel.
	pA = taucs_ccs_create(variablecount, variablecount, 5 * variablecount, TAUCS_DOUBLE);
	pb = new double[3*variablecount];	// All 3 channels at once

	//////////////////////////////////////////////////////////////////////////
	// constraint filed for each channel.
	// all needed channels at once.
	CSimpleImagef confield[3];

	divergence(src, confield[0], 0);
	divergence(src, confield[1], 1);
	divergence(src, confield[2], 2);

	//////////////////////////////////////////////////////////////////////////
	// the buffer pointers.
	double* vals = pA->taucs_values;
	int* rowptr = pA->colptr;
	int* colind = pA->rowind;

	int nequation = 0;
	int nonzerovalueindex = 0;
	
	// Populate matrix
	for (int y = 1; y < height - 1; ++ y)
	{
		for (int x = 1; x < width - 1; ++ x)
		{
			// for each masked pixel.
			if (mask.at(x, y)[0] > 0)
			{
				float channel0 = confield[0].at(x, y)[0];
				float channel1 = confield[1].at(x, y)[0];
				float channel2 = confield[2].at(x, y)[0];

				int pixelindex = y * width + x;
				rowptr[nequation] = nonzerovalueindex;				

				if (mask.at(x, y - 1)[0] > 0)
				{
					vals[nonzerovalueindex] = 1.0;
					colind[nonzerovalueindex] = pixelvariable[pixelindex - width];
					++ nonzerovalueindex;

					//printf("[%d ",pixelvariable[pixelindex - width]);
				}
				else
				{
					// known pixel, update right hand side.
					channel0 -= dest.at(x, y - 1)[0];
					channel1 -= dest.at(x, y - 1)[1];
					channel2 -= dest.at(x, y - 1)[2];
				}

				if (mask.at(x - 1, y)[0] > 0)
				{
					vals[nonzerovalueindex] = 1.0;
					colind[nonzerovalueindex] = pixelvariable[pixelindex - 1];
					++ nonzerovalueindex;

					//printf("%d ",pixelvariable[pixelindex - 1]);
				}
				else
				{
					// known pixel, update right hand side.
					channel0 -= dest.at(x - 1, y)[0];
					channel1 -= dest.at(x - 1, y)[1];
					channel2 -= dest.at(x - 1, y)[2];
				}

				// for pixelself.
				vals[nonzerovalueindex] = -4.0;
				colind[nonzerovalueindex] = pixelvariable[pixelindex];
				++ nonzerovalueindex;

				//////////////////////////////////////////////////////////////////////////
				if (mask.at(x + 1, y)[0] > 0)
				{
					vals[nonzerovalueindex] = 1.0;
					colind[nonzerovalueindex] = pixelvariable[pixelindex + 1];
					++ nonzerovalueindex;

					//printf("%d ",pixelvariable[pixelindex + 1]);
				}
				else
				{
					// known pixel, update right hand side.
					channel0 -= dest.at(x + 1, y)[0];
					channel1 -= dest.at(x + 1, y)[1];
					channel2 -= dest.at(x + 1, y)[2];
				}


				if (mask.at(x, y + 1)[0] > 0)
				{
					vals[nonzerovalueindex] = 1.0;
					colind[nonzerovalueindex] = pixelvariable[pixelindex + width];
					++ nonzerovalueindex;

					//printf("%d] ",pixelvariable[pixelindex + width]);
				}
				else
				{
					// known pixel, update right hand side.
					channel0 -= dest.at(x, y + 1)[0];
					channel1 -= dest.at(x, y + 1)[1];
					channel2 -= dest.at(x, y + 1)[2];
				}


				// Spread the right hand side so we can solve using TAUCS for
				// several channels at once. only update the needed channels.
				pb[nequation] = channel0;
				pb[nequation + variablecount] = channel1;
				pb[nequation + variablecount + variablecount] = channel2;

				//				std::cout<<rr<<","<<gg<<","<<bb<<std::endl;
				++ nequation;
			}
		}
	}

	if (nequation != variablecount)
	{
		std::cout<< "error: parameter error while constructing \n";

		SAFE_DELETE(pA, taucs_ccs_free)
		SAFE_DELETE(pb, delete[])
		*A = NULL;
		*b = NULL;
		return false;
	}
	rowptr[nequation] = nonzerovalueindex; // mark last CRS index

	*A = MatrixTranspose(pA);

	SAFE_DELETE(pA, taucs_ccs_free)
	*b = pb;

	return true;
}

//////////////////////////////////////////////////////////////////////////
/// solve the sparse linear system by taucs solver.
///
/// \param A			 [in] the left matrix of A*x = b.
/// \param b			 [in] the right matrix of A*x = b.
/// \param  channels     [in] how many channels will be rebuilt.
/// \param x			 [out] the solution of A*x = b.
/// \return			     whether the solution is valid.
//////////////////////////////////////////////////////////////////////////
bool solve_discrete_poisson(taucs_ccs_matrix *A, double *b, double ** x)
{
	if (A == NULL || b == NULL)
	{
		return false;
	}

	double *px = (*x);

	SAFE_DELETE(px, delete[])

	// all needed channels at once.
	int channels = 3;

	px = new double[A->n * channels];

	char* options[] = { "taucs.factor.LU=true", NULL }; 

	try
	{
		int error = taucs_linsolve( A, NULL, channels, px, b, options, NULL);
		if (error != TAUCS_SUCCESS)
		{
			std::cout<<"taucs_linsolve failed...."<<error<<std::endl;
			
			SAFE_DELETE(px, delete[])
			*x = NULL;
			return false;
		}	
	}
	catch (...)
	{
		std::cout<<	"taucs_linsolve exception."<<std::endl;
	}

	*x = px;
	return true;
}

//////////////////////////////////////////////////////////////////////////
/// \brief clamp the value to [0, 1.0].
///
/// clamp the pixel value.
///
/// \param pixel   [in] the pixel value.
/// \return        the clamped value.
//////////////////////////////////////////////////////////////////////////
float clamp(float pixel)
{
	if (pixel < 0.0f)
	{
		pixel = 0.0f;
	}
	else if (pixel > 255.0f)
	{
		pixel = 255.0;
	}
	return pixel;
}

void clamp(CSimpleImagef &dest)
{
	for (int y = 0; y < dest._height; ++ y)
	{
		for (int x = 0; x < dest._width; ++ x)
		{
			float *p = dest.at(x, y);

			p[0] = clamp(p[0]);
			p[1] = clamp(p[1]);
			p[2] = clamp(p[2]);
		}
	}
}
//////////////////////////////////////////////////////////////////////////
/// \brief rebuild the image based on the new color.
///
/// copy the pixel one by one. refer to the paper for details.
///
/// \param	tarchannel  [in|out] the target channels which will be rebuilt.
/// \param  maskmap		[in] the image provides the mask.
/// \param  affch       [in] which channels will be rebuilt.
/// \param  C			[in] the array of new color.
/// \return  whether succeed.
//////////////////////////////////////////////////////////////////////////
bool rebuild_image(CSimpleImagef &destimage, CSimpleImageb &mask, double *C)
{
	int height = destimage._height;
	int width = destimage._width;

	int variablecount = 0;
	for (int y = 1; y < height - 1; ++ y)
	{
		for (int x = 1; x < width - 1; ++ x)
		{
			if (mask.at(x, y)[0] > 0)
			{
				++ variablecount;
			}
		}
	}
	
	int index = 0;
	for (int y = 1; y < height - 1; ++ y)
	{
		for (int x = 1; x < width - 1; ++ x)
		{
			if (mask.at(x, y)[0] > 0)
			{
				destimage.at(x, y)[0] = (C[index]);
				destimage.at(x, y)[1] = (C[index + variablecount]);
				destimage.at(x, y)[2] = (C[index + variablecount + variablecount]);

				++ index;
			}
		}
	}

	return true;
}

void main(int argc, char** argv)
{
	if (argc < 7)
	{
		printf("usage: PossionEditing.exe src_image src_mask dest_image offsetx offsety result_image\n");
		return;
	}

	std::string srcfile = argv[1];
	std::string srcmask = argv[2];

	std::string destfile = argv[3];

	int offsetx = atoi(argv[4]);
	int offsety = atoi(argv[5]);

	std::string result = argv[6];

	CSimpleImagef srcimage;
	CSimpleImageb srcmaskimage;

	CSimpleImagef destimage;

	CxImage image;


	SimpleImage::load2ximage(srcimage, srcfile);
	SimpleImage::load2ximage(srcmaskimage, srcmask);
	SimpleImage::load2ximage(destimage, destfile);

	if (srcimage._width != srcmaskimage._width || srcimage._height != srcmaskimage._height)
	{
		return;
	}

	SimpleImage::flip(srcimage);
	SimpleImage::flip(destimage);
	SimpleImage::flip(srcmaskimage);

	//////////////////////////////////////////////////////////////////////////
	CSimpleImagef newsrcimage = destimage;

	CSimpleImageb newmaskimage;

	newmaskimage.set_size(destimage._width, destimage._height, 1);
	SimpleImage::zero(newmaskimage);

	//////////////////////////////////////////////////////////////////////////
	for (int y = 0; y < srcmaskimage._height; ++ y)
	{
		for (int x = 0; x < srcmaskimage._width; ++ x)
		{
			if (x + offsetx >= 0 && x + offsetx < destimage._width &&
				y + offsety >= 0 && y + offsety < destimage._height)
			{
				float *psrc = srcimage.at(x, y);
				float *pdest = newsrcimage.at(x + offsetx, y + offsety);

				pdest[0] = psrc[0];
				pdest[1] = psrc[1];
				pdest[2] = psrc[2];

				newmaskimage.at(x + offsetx, y + offsety)[0] = srcmaskimage.at(x, y)[0];
			}
		}
	}

	//SimpleImage::save2ximage(newsrcimage, "d:/poisson/newimage.jpg");
	//SimpleImage::save2ximage(newmaskimage, "d:/poisson/newmask.jpg");

	taucs_ccs_matrix* A = NULL;
	double *b = NULL;
	double *C = NULL;

	//RGBtoYUV(newsrcimage);
	//RGBtoYUV(destimage);

	//CSimpleImagef save_src = newsrcimage;

	if (build_discrete_poisson(newsrcimage, destimage, newmaskimage, &A, &b))
	{
		if (solve_discrete_poisson(A, b, &C))
		{
			rebuild_image(destimage, newmaskimage, C);
		}
	}

	//for (int y = 0; y < destimage._height; ++ y)
	//{
	//	for (int x = 0; x < destimage._width; ++ x)
	//	{
	//		if (newmaskimage.at(x, y)[0] > 0)
	//		{
	//			float *p0 = destimage.at(x, y);
	//			float *p1 = save_src.at(x, y);

	//			p0[1] = p1[1];
	//			p0[2] = p1[2];
	//		}
	//	}
	//}

	//YUVtoRGB(destimage);

	clamp(destimage);

	SimpleImage::flip(destimage);

	SimpleImage::save2ximage(destimage, result);
}