/* solve and invert matrices
 *
 * 19/4/20 kleisauke
 *	- from im_matinv
 */

/*

	This file is part of VIPS.

	VIPS is free software; you can redistribute it and/or modify
	it under the terms of the GNU Lesser General Public License as published by
	the Free Software Foundation; either version 2 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 Lesser General Public License for more details.

	You should have received a copy of the GNU Lesser General Public License
	along with this program; if not, write to the Free Software
	Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
	02110-1301  USA

 */

/*

	These files are distributed with VIPS - http://www.vips.ecs.soton.ac.uk

 */

/*
#define DEBUG
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif /*HAVE_CONFIG_H*/
#include <glib/gi18n-lib.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

#include <vips/vips.h>

/* Our state.
 */
typedef struct _VipsMatrixinvert {
	VipsOperation parent_instance;

	VipsImage *in;
	VipsImage *out;

	/* .. and cast to a matrix.
	 */
	VipsImage *mat;

	/* The LU decomposed matrix.
	 */
	VipsImage *lu;
} VipsMatrixinvert;

typedef VipsOperationClass VipsMatrixinvertClass;

G_DEFINE_TYPE(VipsMatrixinvert, vips_matrixinvert, VIPS_TYPE_OPERATION);

static void
vips_matrixinvert_dispose(GObject *gobject)
{
	VipsMatrixinvert *matrix = (VipsMatrixinvert *) gobject;

	VIPS_UNREF(matrix->mat);
	VIPS_UNREF(matrix->lu);

	G_OBJECT_CLASS(vips_matrixinvert_parent_class)->dispose(gobject);
}

/* DBL_MIN is smallest *normalized* double precision float
 */
#define TOO_SMALL (2.0 * DBL_MIN)

/* Save a bit of typing.
 */
#define ME(m, i, j) (*VIPS_MATRIX((m), (i), (j)))

/**
 * lu_decomp:
 * @mat: matrix to decompose
 *
 * This function takes any square NxN [class@Image].
 * It returns a [class@Image] which is (N+1)xN.
 *
 * It calculates the PLU decomposition, storing the upper and diagonal parts
 * of U, together with the lower parts of L, as an NxN matrix in the first
 * N rows of the new matrix.  The diagonal parts of L are all set to unity
 * and are not stored.
 *
 * The final row of the new [class@Image] has only integer entries, which
 * represent the row-wise permutations made by the permutation matrix P.
 *
 * The scale and offset members of the input [class@Image] are ignored.
 *
 * See:
 *
 *   PRESS, W. et al, 1992.  Numerical Recipes in C; The Art of Scientific
 *   Computing, 2nd ed.  Cambridge: Cambridge University Press, pp. 43-50.
 *
 * Returns: the decomposed matrix on success, or `NULL` on error.
 */
static VipsImage *
lu_decomp(VipsImage *mat)
{
	int i, j, k;
	double *row_scale;
	VipsImage *lu;

	if (!(row_scale = VIPS_ARRAY(NULL, mat->Xsize, double))) {
		return NULL;
	}

	if (!(lu = vips_image_new_matrix(mat->Xsize, mat->Xsize + 1))) {
		g_free(row_scale);
		return NULL;
	}

	/* copy all coefficients and then perform decomposition in-place */
	memcpy(VIPS_MATRIX(lu, 0, 0), VIPS_MATRIX(mat, 0, 0),
		(size_t) mat->Xsize * mat->Xsize * sizeof(double));

	for (i = 0; i < mat->Xsize; ++i) {
		row_scale[i] = 0.0;

		for (j = 0; j < mat->Xsize; ++j) {
			double abs_val = fabs(ME(lu, i, j));

			/* find largest in each ROW */
			if (abs_val > row_scale[i])
				row_scale[i] = abs_val;
		}

		if (!row_scale[i]) {
			vips_error("matrixinvert", "singular matrix");
			g_object_unref(lu);
			g_free(row_scale);
			return NULL;
		}

		/* fill array with scaling factors for each ROW */
		row_scale[i] = 1.0 / row_scale[i];
	}

	for (j = 0; j < mat->Xsize; ++j) { /* loop over COLs */
		double max = -1.0;
		int i_of_max;

		/* not needed, but stops a compiler warning */
		i_of_max = 0;

		/* loop over ROWS in upper-half, except diagonal */
		for (i = 0; i < j; ++i)
			for (k = 0; k < i; ++k)
				ME(lu, i, j) -= ME(lu, i, k) * ME(lu, k, j);

		/* loop over ROWS in diagonal and lower-half */
		for (i = j; i < mat->Xsize; ++i) {
			double abs_val;

			for (k = 0; k < j; ++k)
				ME(lu, i, j) -= ME(lu, i, k) * ME(lu, k, j);

			/* find largest element in each COLUMN scaled so that */
			/* largest in each ROW is 1.0 */
			abs_val = row_scale[i] * fabs(ME(lu, i, j));

			if (abs_val > max) {
				max = abs_val;
				i_of_max = i;
			}
		}

		if (fabs(ME(lu, i_of_max, j)) < TOO_SMALL) {
			/* divisor is near zero */
			vips_error("matrixinvert", "singular or near-singular matrix");
			g_object_unref(lu);
			g_free(row_scale);
			return NULL;
		}

		if (i_of_max != j) {
			/* swap ROWS */
			for (k = 0; k < mat->Xsize; ++k) {
				double temp = ME(lu, j, k);
				ME(lu, j, k) = ME(lu, i_of_max, k);
				ME(lu, i_of_max, k) = temp;
			}

			row_scale[i_of_max] = row_scale[j];
			/* no need to copy this scale back up - we won't use it */
		}

		/* record permutation */
		ME(lu, j, mat->Xsize) = i_of_max;

		/* divide by best (largest scaled) pivot found */
		for (i = j + 1; i < mat->Xsize; ++i)
			ME(lu, i, j) /= ME(lu, j, j);
	}
	g_free(row_scale);

	return lu;
}

/**
 * lu_solve:
 * @lu: matrix to solve
 * @vec: name for output matrix
 *
 * Solve the system of linear equations Ax=b, where matrix A has already
 * been decomposed into LU form in @lu.  Input vector b is in
 * vec and is overwritten with vector x.
 *
 * See:
 *
 *   PRESS, W. et al, 1992.  Numerical Recipes in C; The Art of Scientific
 *   Computing, 2nd ed.  Cambridge: Cambridge University Press, pp. 43-50.
 *
 * ::: seealso
 *     [method@Image.matrixmultiply].
 *
 * Returns: 0 on success, -1 on error
 */
static int
lu_solve(VipsImage *lu, double *vec)
{
	int i, j;

	if (lu->Xsize + 1 != lu->Ysize) {
		vips_error("matrixinvert", "not an LU decomposed matrix");
		return -1;
	}

	for (i = 0; i < lu->Xsize; ++i) {
		int i_perm = ME(lu, i, lu->Xsize);

		if (i_perm != i) {
			double temp = vec[i];
			vec[i] = vec[i_perm];
			vec[i_perm] = temp;
		}
		for (j = 0; j < i; ++j)
			vec[i] -= ME(lu, i, j) * vec[j];
	}

	for (i = lu->Xsize - 1; i >= 0; --i) {

		for (j = i + 1; j < lu->Xsize; ++j)
			vec[i] -= ME(lu, i, j) * vec[j];

		vec[i] /= ME(lu, i, i);
	}

	return 0;
}

static int
vips_matrixinvert_solve(VipsMatrixinvert *matrix)
{
	VipsImage *out = matrix->out;

	int i, j;
	double *vec;

	if (!(matrix->lu = lu_decomp(matrix->mat)))
		return -1;

	if (!(vec = VIPS_ARRAY(matrix, matrix->lu->Xsize, double)))
		return -1;

	for (j = 0; j < matrix->lu->Xsize; ++j) {
		for (i = 0; i < matrix->lu->Xsize; ++i)
			vec[i] = 0.0;

		vec[j] = 1.0;

		if (lu_solve(matrix->lu, vec))
			return -1;

		for (i = 0; i < matrix->lu->Xsize; ++i)
			ME(out, i, j) = vec[i];
	}

	return 0;
}

static int
vips_matrixinvert_direct(VipsMatrixinvert *matrix)
{
	VipsObjectClass *class = VIPS_OBJECT_GET_CLASS(matrix);
	VipsImage *in = matrix->mat;
	VipsImage *out = matrix->out;

	switch (matrix->mat->Xsize) {
	case 1: {
		double det = ME(in, 0, 0);

		if (fabs(det) < TOO_SMALL) {
			/* divisor is near zero */
			vips_error(class->nickname,
				"%s", _("singular or near-singular matrix"));
			return -1;
		}

		ME(out, 0, 0) = 1.0 / det;
	} break;

	case 2: {
		double det = ME(in, 0, 0) * ME(in, 1, 1) -
			ME(in, 0, 1) * ME(in, 1, 0);

		double tmp;

		if (fabs(det) < TOO_SMALL) {
			/* divisor is near zero */
			vips_error(class->nickname,
				"%s", _("singular or near-singular matrix"));
			return -1;
		}

		tmp = 1.0 / det;
		ME(out, 0, 0) = tmp * ME(in, 1, 1);
		ME(out, 0, 1) = -tmp * ME(in, 0, 1);
		ME(out, 1, 0) = -tmp * ME(in, 1, 0);
		ME(out, 1, 1) = tmp * ME(in, 0, 0);
	} break;

	case 3: {
		double det;
		double tmp;

		det = ME(in, 0, 0) *
			(ME(in, 1, 1) * ME(in, 2, 2) - ME(in, 1, 2) * ME(in, 2, 1));
		det -= ME(in, 0, 1) *
			(ME(in, 1, 0) * ME(in, 2, 2) - ME(in, 1, 2) * ME(in, 2, 0));
		det += ME(in, 0, 2) *
			(ME(in, 1, 0) * ME(in, 2, 1) - ME(in, 1, 1) * ME(in, 2, 0));

		if (fabs(det) < TOO_SMALL) {
			/* divisor is near zero */
			vips_error(class->nickname,
				"%s", _("singular or near-singular matrix"));
			return -1;
		}

		tmp = 1.0 / det;

		ME(out, 0, 0) = tmp *
			(ME(in, 1, 1) * ME(in, 2, 2) - ME(in, 1, 2) * ME(in, 2, 1));
		ME(out, 1, 0) = tmp *
			(ME(in, 1, 2) * ME(in, 2, 0) - ME(in, 1, 0) * ME(in, 2, 2));
		ME(out, 2, 0) = tmp *
			(ME(in, 1, 0) * ME(in, 2, 1) - ME(in, 1, 1) * ME(in, 2, 0));

		ME(out, 0, 1) = tmp *
			(ME(in, 0, 2) * ME(in, 2, 1) - ME(in, 0, 1) * ME(in, 2, 2));
		ME(out, 1, 1) = tmp *
			(ME(in, 0, 0) * ME(in, 2, 2) - ME(in, 0, 2) * ME(in, 2, 0));
		ME(out, 2, 1) = tmp *
			(ME(in, 0, 1) * ME(in, 2, 0) - ME(in, 0, 0) * ME(in, 2, 1));

		ME(out, 0, 2) = tmp *
			(ME(in, 0, 1) * ME(in, 1, 2) - ME(in, 0, 2) * ME(in, 1, 1));
		ME(out, 1, 2) = tmp *
			(ME(in, 0, 2) * ME(in, 1, 0) - ME(in, 0, 0) * ME(in, 1, 2));
		ME(out, 2, 2) = tmp *
			(ME(in, 0, 0) * ME(in, 1, 1) - ME(in, 0, 1) * ME(in, 1, 0));
	} break;

	/* TODO(kleisauke):
	 * We sometimes use 4x4 matrices, could we also make a
	 * direct version for those? For e.g.:
	 * https://stackoverflow.com/a/1148405/10952119 */
	default:
		g_assert(0);
		return -1;
	}

	return 0;
}

static int
vips_matrixinvert_build(VipsObject *object)
{
	VipsObjectClass *class = VIPS_OBJECT_GET_CLASS(object);
	VipsMatrixinvert *matrix = (VipsMatrixinvert *) object;

	if (VIPS_OBJECT_CLASS(vips_matrixinvert_parent_class)->build(object))
		return -1;

	if (vips_check_matrix(class->nickname, matrix->in, &matrix->mat))
		return -1;

	if (matrix->mat->Xsize != matrix->mat->Ysize) {
		vips_error(class->nickname, "%s", _("non-square matrix"));
		return -1;
	}

	g_object_set(matrix,
		"out", vips_image_new_matrix(matrix->mat->Xsize, matrix->mat->Ysize),
		NULL);

	/* Direct path for < 4x4 matrices
	 */
	if (matrix->mat->Xsize >= 4) {
		if (vips_matrixinvert_solve(matrix))
			return -1;
	}
	else {
		if (vips_matrixinvert_direct(matrix))
			return -1;
	}

	return 0;
}

static void
vips_matrixinvert_class_init(VipsMatrixinvertClass *class)
{
	GObjectClass *gobject_class = G_OBJECT_CLASS(class);
	VipsObjectClass *vobject_class = VIPS_OBJECT_CLASS(class);

	gobject_class->dispose = vips_matrixinvert_dispose;
	gobject_class->set_property = vips_object_set_property;
	gobject_class->get_property = vips_object_get_property;

	vobject_class->nickname = "matrixinvert";
	vobject_class->description = _("invert a matrix");
	vobject_class->build = vips_matrixinvert_build;

	VIPS_ARG_IMAGE(class, "in", 0,
		_("Input"),
		_("An square matrix"),
		VIPS_ARGUMENT_REQUIRED_INPUT,
		G_STRUCT_OFFSET(VipsMatrixinvert, in));

	VIPS_ARG_IMAGE(class, "out", 1,
		_("Output"),
		_("Output matrix"),
		VIPS_ARGUMENT_REQUIRED_OUTPUT,
		G_STRUCT_OFFSET(VipsMatrixinvert, out));
}

static void
vips_matrixinvert_init(VipsMatrixinvert *matrix)
{
}

/**
 * vips_matrixinvert: (method)
 * @m: matrix to invert
 * @out: (out): output matrix
 * @...: `NULL`-terminated list of optional named arguments
 *
 * This operation calculates the inverse of the matrix represented in @m.
 * The scale and offset members of the input matrix are ignored.
 *
 * ::: seealso
 *     [ctor@Image.matrixload], [method@Image.matrixmultiply].
 *
 * Returns: 0 on success, -1 on error
 */
int
vips_matrixinvert(VipsImage *m, VipsImage **out, ...)
{
	va_list ap;
	int result;

	va_start(ap, out);
	result = vips_call_split("matrixinvert", ap, m, out);
	va_end(ap);

	return result;
}
