/*
    Scan Tailor - Interactive post-processing tool for scanned pages.
    Copyright (C) 2015  Joseph Artsimovich <joseph.artsimovich@gmail.com>

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "Optimizer.h"
#include <Eigen/Core>
#include <Eigen/QR>
#include <boost/foreach.hpp>
#include <stdexcept>
#include <algorithm>
#include <cassert>

using namespace Eigen;

namespace spfit
{

Optimizer::Optimizer(size_t num_vars)
:	m_numVars(num_vars)
,	m_A(num_vars, num_vars)
,	m_b(num_vars)
,	m_x(num_vars)
,	m_externalForce(num_vars)
,	m_internalForce(num_vars)
{
	m_A.setZero();
	m_b.setZero();
	m_x.setZero();
}

void
Optimizer::setConstraints(std::list<LinearFunction> const& constraints)
{
	size_t const num_constraints = constraints.size();
	size_t const num_dimensions = m_numVars + num_constraints;

	MatrixXd A(num_dimensions, num_dimensions);
	A.setZero();
	VectorXd b(num_dimensions);
	b.setZero();
	// Matrix A and vector b will have the following layout:
	//     |N N N L L|      |-D|
	//     |N N N L L|      |-D|
	// A = |N N N L L|  b = |-D|
	//     |C C C 0 0|      |-J|
	//     |C C C 0 0|      |-J|
	// N: non-constant part of the gradient of the function we are minimizing.
	// C: non-constant part of constraint functions (one per line).
	// L: coefficients of Lagrange multipliers.  These happen to be equal
	//    to the symmetric C values.
	// D: constant part of the gradient of the function we are optimizing.
	// J: constant part of constraint functions.

	std::list<LinearFunction>::const_iterator ctr(constraints.begin());
	for (size_t i = m_numVars; i < num_dimensions; ++i, ++ctr) {
		b[i] = -ctr->b;
		for (size_t j = 0; j < m_numVars; ++j) {
			A(i, j) = A(j, i) = ctr->a[j];
		}
	}

	m_x.resize(num_dimensions);
	m_x.setZero();
	m_A.swap(A);
	m_b.swap(b);
}

void
Optimizer::addExternalForce(QuadraticFunction const& force)
{
	m_externalForce += force;
}

void
Optimizer::addExternalForce(QuadraticFunction const& force, std::vector<int> const& sparse_map)
{
	size_t const num_vars = force.numVars();
	for (size_t i = 0; i < num_vars; ++i) {
		int const ii = sparse_map[i];
		for (size_t j = 0; j < num_vars; ++j) {
			int const jj = sparse_map[j];
			m_externalForce.A(ii, jj) += force.A(i, j);
		}
		m_externalForce.b[ii] += force.b[i];
	}
	m_externalForce.c += force.c;
}

void
Optimizer::addInternalForce(QuadraticFunction const& force)
{
	m_internalForce += force;
}

void
Optimizer::addInternalForce(
	QuadraticFunction const& force, std::vector<int> const& sparse_map)
{
	size_t const num_vars = force.numVars();
	for (size_t i = 0; i < num_vars; ++i) {
		int const ii = sparse_map[i];
		for (size_t j = 0; j < num_vars; ++j) {
			int const jj = sparse_map[j];
			m_internalForce.A(ii, jj) += force.A(i, j);
		}
		m_internalForce.b[ii] += force.b[i];
	}
	m_internalForce.c += force.c;
}

OptimizationResult
Optimizer::optimize(double internal_force_weight)
{
	// Note: because we are supposed to reset the forces anyway,
	// we re-use m_internalForce to store the cummulative force.
	m_internalForce *= internal_force_weight;
	m_internalForce += m_externalForce;
	
	// For the layout of m_A and m_b, see setConstraints()
	QuadraticFunction::Gradient const grad(m_internalForce.gradient());
	for (size_t i = 0; i < m_numVars; ++i) {
		m_b[i] = -grad.b[i];
		for (size_t j = 0; j < m_numVars; ++j) {
			m_A(i, j) = grad.A(i, j);
		}
	}

	double const total_force_before = m_internalForce.c;

	auto qr = m_A.colPivHouseholderQr();
	if (!qr.isInvertible()) {
		m_externalForce.reset();
		m_internalForce.reset();
		m_x.setZero(); // To make undoLastStep() work as expected.
		return OptimizationResult(total_force_before, total_force_before);
	}

	m_x = qr.solve(m_b).head(m_numVars);
	double const total_force_after = m_internalForce.evaluate(m_x);
	m_externalForce.reset(); // Now it's finally safe to reset these.
	m_internalForce.reset();

	// The last thing remaining is to adjust constraints,
	// as they depend on the current variables.
	adjustConstraints(1.0);

	return OptimizationResult(total_force_before, total_force_after);
}

void
Optimizer::undoLastStep()
{
	adjustConstraints(-1.0);
	m_x.fill(0);
}

/**
 * direction == 1 is used for forward adjustment,
 * direction == -1 is used for undoing the last step.
 */
void
Optimizer::adjustConstraints(double direction)
{
	size_t const num_dimensions = m_b.size();
	for (size_t i = m_numVars; i < num_dimensions; ++i) {
		// See setConstraints() for more information
		// on the layout of m_A and m_b.
		double c = 0;
		for (size_t j = 0; j < m_numVars; ++j) {
			c += m_A(i, j) * m_x[j];
		}
		m_b[i] -= c * direction;
	}
}

void
Optimizer::swap(Optimizer& other)
{
	m_A.swap(other.m_A);
	m_b.swap(other.m_b);
	m_x.swap(other.m_x);
	m_externalForce.swap(other.m_externalForce);
	m_internalForce.swap(other.m_internalForce);
	std::swap(m_numVars, other.m_numVars);
}

} // namespace spfit
