#include <AOFlagger/strategy/algorithms/polfitmethod.h>

PolFitMethod::PolFitMethod() : _background(0), _previousCoefficients(0)
{
}

PolFitMethod::~PolFitMethod()
{
	if(_background!=0)
		delete _background;
}

void PolFitMethod::Initialize(const TimeFrequencyData &input)
{
	_original = input.GetSingleImage();
	_background2D = Image2D::CreateEmptyImagePtr(_original->Width(), _original->Height());
	if(_background!=0)
		delete _background;
	_background = new TimeFrequencyData(input.PhaseRepresentation(), input.Polarisation(), _background2D);
	_mask = input.GetSingleMask();
	if(_hSquareSize * 2 > _original->Width())
		_hSquareSize = _original->Width()/2;
	if(_vSquareSize * 2 > _original->Height())
		_vSquareSize = _original->Height()/2;
}

unsigned PolFitMethod::TaskCount()
{
	return _original->Height();
}

void PolFitMethod::PerformFit(unsigned taskNumber)
{
	if(_mask == 0)
		throw BadUsageException("Mask has not been set!");
	unsigned y = taskNumber;
	for(unsigned x=0;x<_original->Width();++x)
		_background2D->SetValue(x, y, CalculateBackgroundValue(x, y));
}

long double PolFitMethod::CalculateBackgroundValue(unsigned x, unsigned y)
{
	ThreadLocal local;
	local.image = this;
	local.currentX = x;
	local.currentY = y;

	if(local.currentY >= _vSquareSize)
		local.startY = local.currentY-_vSquareSize;
	else
		local.startY = 0;
	local.endY = local.currentY + _vSquareSize;
	if(local.endY >= _original->Height())
		local.endY = _original->Height()-1;

	if(local.currentX >= _hSquareSize)
		local.startX = local.currentX - _hSquareSize;
	else
		local.startX = 0;
	local.endX = local.currentX + _hSquareSize;
		if(local.endX >= _original->Width())
			local.endX = _original->Width()-1;
	local.emptyWindows = 0;

	switch(_method) {
		case None:
			return 0.0;
		case LeastSquare:
		case LeastAbs:
		case FastGaussianWeightedAverage:
			return FitBackground(x, y, local);
		default:
			throw ConfigurationException("The PolFitMethod was not initialized before a fit was executed.");
	}
}

int PolFitMethod::SquareError(const gsl_vector * coefs, void *data, gsl_vector * f)
{
	// f(x,y) = ( a * x^2 + b * xy + c * y^2 + d x + e y + l  -  image[x,y] )^2
	ThreadLocal &local = *(ThreadLocal *) data;
	double a = gsl_vector_get(coefs, 0);
	double b = gsl_vector_get(coefs, 1);
	double c = gsl_vector_get(coefs, 2);
	double d = gsl_vector_get(coefs, 3);
	double e = gsl_vector_get(coefs, 4);
	double l = gsl_vector_get(coefs, 5);

	unsigned index = 0;

	for(unsigned y=local.startY; y <= local.endY; ++y) {
		double yf = y;
		for(unsigned x=local.startX; x <= local.endX; ++x) {
			if(!local.image->_mask->Value(x, y)) {
				double xf = x;
				double g_xy = a * xf*xf + b * xf*yf + c * yf*yf + d * xf + e * yf + l - local.image->_original->Value(x, y);
				gsl_vector_set(f, index, (g_xy*g_xy));
			} else {
				gsl_vector_set(f, index, 0.0);
			}
			index++;
		}
	}
	return GSL_SUCCESS;
}

int PolFitMethod::SquareErrorDiff(const gsl_vector *coefs, void *data, gsl_matrix * J)
{
	ThreadLocal &local = *(ThreadLocal *) data;

	double a = gsl_vector_get(coefs, 0);
	double b = gsl_vector_get(coefs, 1);
	double c = gsl_vector_get(coefs, 2);
	double d = gsl_vector_get(coefs, 3);
	double e = gsl_vector_get(coefs, 4);
	double f = gsl_vector_get(coefs, 5);

	unsigned index = 0;
	for(unsigned y=local.startY; y <= local.endY; ++y) {
		double yf = y;
		for(unsigned x=local.startX; x <= local.endX; ++x) {
			if(!local.image->_mask->Value(x, y)) {
				// f(x,y) = ( a * x^2 + b * xy + c * y^2 + d x + e y + f  -  image[x,y] )^2
				// f(x,y) = g^2(x,y)
				// df(x,y)/dz =  2 * g(x,y) * dg(x,y)/dz
				// We now need to calculate df(x,y)/da, df(x,y)/db, ..... for each x and y
				double xf = x;
				double g_xy = 2.0 * (a * xf*xf + b * xf*yf + c * yf*yf + d * xf + e * yf + f - local.image->_original->Value(x, y));
				gsl_matrix_set(J, index, 0, g_xy * xf*xf); // df/da
				gsl_matrix_set(J, index, 1, g_xy * xf*yf); // df/db
				gsl_matrix_set(J, index, 2, g_xy * yf*yf); // df/dc
				gsl_matrix_set(J, index, 3, g_xy * xf); // df/dd
				gsl_matrix_set(J, index, 4, g_xy * yf); // df/de
				gsl_matrix_set(J, index, 5, g_xy); // df/df
			} else {
				for(unsigned ci=0;ci<6;++ci)
					gsl_matrix_set(J, index, ci, 0.0);
			}
			index++;
		}
	}
	
	return GSL_SUCCESS;
}

int PolFitMethod::LinError(const gsl_vector * coefs, void *data, gsl_vector * f)
{
	// f(x,y) = | a * x^2 + b * xy + c * y^2 + d x + e y + l  -  image[x,y] |
	ThreadLocal &local = *(ThreadLocal *) data;
	double a = gsl_vector_get(coefs, 0);
	double b = gsl_vector_get(coefs, 1);
	double c = gsl_vector_get(coefs, 2);
	double d = gsl_vector_get(coefs, 3);
	double e = gsl_vector_get(coefs, 4);
	double l = gsl_vector_get(coefs, 5);

	unsigned index = 0;

	for(unsigned y=local.startY; y <= local.endY; ++y) {
		double yf = y;
		for(unsigned x=local.startX; x <= local.endX; ++x) {
			if(!local.image->_mask->Value(x, y)) {
				double xf = x;
				double g_xy = a * xf*xf + b * xf*yf + c * yf*yf + d * xf + e * yf + l - local.image->_original->Value(x, y);
				gsl_vector_set(f, index, fabs(g_xy));
			} else {
				gsl_vector_set(f, index, 0.0);
			}
			index++;
		}
	}
	return GSL_SUCCESS;
}

int PolFitMethod::LinErrorDiff(const gsl_vector *coefs, void *data, gsl_matrix * J)
{
	ThreadLocal &local = *(ThreadLocal *) data;

	double a = gsl_vector_get(coefs, 0);
	double b = gsl_vector_get(coefs, 1);
	double c = gsl_vector_get(coefs, 2);
	double d = gsl_vector_get(coefs, 3);
	double e = gsl_vector_get(coefs, 4);
	double f = gsl_vector_get(coefs, 5);

	unsigned index = 0;
	for(unsigned y=local.startY; y <= local.endY; ++y) {
		double yf = y;
		for(unsigned x=local.startX; x <= local.endX; ++x) {
			if(!local.image->_mask->Value(x, y)) {
				// f(x,y) = | a * x^2 + b * xy + c * y^2 + d x + e y + f  -  image[x,y] |
				// f(x,y) = | g(x,y) |
				// df(x,y)/dz =  dg(x,y)/dz * g(x,y) / | g(x,y) |
				// We now need to calculate df(x,y)/da, df(x,y)/db, ..... for each x and y
				double xf = x;
				double g_xy = a * xf*xf + b * xf*yf + c * yf*yf + d * xf + e * yf + f - local.image->_original->Value(x, y);
				double h_xy = g_xy / fabs(g_xy);
				gsl_matrix_set(J, index, 0, h_xy * xf*xf); // df/da
				gsl_matrix_set(J, index, 1, h_xy * xf*yf); // df/db
				gsl_matrix_set(J, index, 2, h_xy * yf*yf); // df/dc
				gsl_matrix_set(J, index, 3, h_xy * xf); // df/dd
				gsl_matrix_set(J, index, 4, h_xy * yf); // df/de
				gsl_matrix_set(J, index, 5, h_xy); // df/df
			} else {
				for(unsigned ci=0;ci<6;++ci)
					gsl_matrix_set(J, index, ci, 0.0);
			}
			index++;
		}
	}
	
	return GSL_SUCCESS;
}

long double PolFitMethod::FitBackground(unsigned x, unsigned y, ThreadLocal &local)
{
	std::vector<long double> coefficients(6);

	boost::mutex::scoped_lock lock(_mutex);
	if(_previousCoefficients) {
		for(unsigned i=0;i<6;++i)
			coefficients[i] = _previousCoefficients[i];
	} else {
		for(unsigned i=0;i<6;++i)
			coefficients[i] = 1e-4 * (i*i*i);
	}
	lock.unlock();

	// Chose to use the Levenberg-Marquardt solver with scaling
	const gsl_multifit_fdfsolver_type * T = gsl_multifit_fdfsolver_lmsder;

	// Construct solver
	unsigned functionCount = (local.endY - local.startY + 1) * (local.endX - local.startX + 1);
	unsigned coefficientCount = 6;
	gsl_multifit_fdfsolver *solver = gsl_multifit_fdfsolver_alloc (T, functionCount, coefficientCount);
	if(solver == 0) throw BadUsageException("No solver.");

	// Initialize function information structure
	gsl_multifit_function_fdf functionInfo;
	switch(_method) {
		case LeastSquare:
		default:
			functionInfo.f = &SquareError;
			functionInfo.df = &SquareErrorDiff;
			functionInfo.fdf = &SquareErrorComb;
			break;
		case LeastAbs:
			functionInfo.f = &LinError;
			functionInfo.df = &LinErrorDiff;
			functionInfo.fdf = &LinErrorComb;
			break;
	}
	functionInfo.n = functionCount;
	functionInfo.p = coefficientCount;
	functionInfo.params = &local;

	// Initialize initial value of parameters
	//gsl_vector vec;
	double vec_init[coefficientCount];
	for(unsigned i = 0;i < coefficientCount;++i)
		vec_init[i] = coefficients[i];
	gsl_vector_view vec_view = gsl_vector_view_array (vec_init, coefficientCount);
	
	int status = gsl_multifit_fdfsolver_set (solver, &functionInfo, &vec_view.vector);
	if (status && status != GSL_CONTINUE) {
		std::cout << "Error: status = " << gsl_strerror (status) << std::endl;
	}

	// Start iterating
	int iter=0;
	do {
		iter++;
		status = gsl_multifit_fdfsolver_iterate(solver);
		//PrintState(iter, solver);

		if (status && status != GSL_CONTINUE) {
			//std::cout << "Error: status = " << gsl_strerror (status) << std::endl;
			break;
		}

		status = gsl_multifit_test_delta(solver->dx, solver->x, _precision, _precision);
	} while (status == GSL_CONTINUE && iter < 250);
	
	// Save coefficients
	for(unsigned i = 0;i<coefficientCount;++i)
		coefficients[i] = gsl_vector_get(solver->x, i);

	lock.lock();
	if(_previousCoefficients == 0)
		_previousCoefficients = new long double[6];
	for(unsigned i = 0;i<coefficientCount;++i)
		_previousCoefficients[i] = coefficients[i];
	lock.unlock();

	// Clean up
	gsl_multifit_fdfsolver_free(solver);

	long double evaluation = Evaluate(x, y, coefficients);
	return evaluation;
}

long double PolFitMethod::Evaluate(unsigned x, unsigned y, long double *coefficients)
{
	// f(x,y) = a * x^2 + b * xy + c * y^2 + d x + e y + f
	double xf = x, yf = y;
	return coefficients[0]*xf*xf + coefficients[1]*xf*yf + coefficients[2]*yf*yf + coefficients[3]*xf +
		coefficients[4]*yf + coefficients[5];
}
