/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2023, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

#include "math-precomp.h"  // Precompiled headers
//
#include <mrpt/math/CSparseMatrix.h>

#include <cstring>
#include <iostream>
#include <string>

using std::cout;
using std::endl;
using std::string;
using namespace mrpt;
using namespace mrpt::math;

/** Copy constructor */
CSparseMatrix::CSparseMatrix(const CSparseMatrix& other)
{
	construct_from_existing_cs(other.sparse_matrix);
	copy(&other.sparse_matrix);
}

/** Copy constructor from an existing "cs" CSparse data structure */
CSparseMatrix::CSparseMatrix(const cs* const sm)
{
	construct_from_existing_cs(*sm);
	copy(sm);
}

/** Copy the data from an existing "cs" CSparse data structure */
void CSparseMatrix::copy(const cs* const sm)
{
	ASSERTMSG_(
		sm->nz == -1,
		"I expected a column-compressed sparse matrix, not a triplet form.");

	sparse_matrix.m = sm->m;
	sparse_matrix.n = sm->n;
	sparse_matrix.nz = sm->nz;
	sparse_matrix.nzmax = sm->nzmax;

	// Do it in a different way depending on it being triplet or compressed
	// format:
	std::memcpy(sparse_matrix.i, sm->i, sizeof(int) * sm->nzmax);
	std::memcpy(sparse_matrix.p, sm->p, sizeof(int) * (sm->n + 1));
	std::memcpy(sparse_matrix.x, sm->x, sizeof(double) * sm->nzmax);
}

/** Fast copy the data from an existing "cs" CSparse data structure, copying the
 * pointers and leaving NULLs in the source structure. */
void CSparseMatrix::copy_fast(cs* const sm)
{
	// This method will work for either triplet or compressed format:

	// Free previous contents, if any.
	internal_free_mem();

	// Fast copy / Move:
	sparse_matrix.m = sm->m;
	sparse_matrix.n = sm->n;
	sparse_matrix.nz = sm->nz;
	sparse_matrix.nzmax = sm->nzmax;

	sparse_matrix.i = sm->i;
	sparse_matrix.p = sm->p;
	sparse_matrix.x = sm->x;

	// Mark source as empty:
	sm->i = nullptr;
	sm->p = nullptr;
	sm->x = nullptr;
}

/** Fast swap contents with another sparse matrix */
void CSparseMatrix::swap(CSparseMatrix& other)
{
	// Fast copy / Move:
	std::swap(sparse_matrix.m, other.sparse_matrix.m);
	std::swap(sparse_matrix.n, sparse_matrix.n);
	std::swap(sparse_matrix.nz, other.sparse_matrix.nz);
	std::swap(sparse_matrix.nzmax, other.sparse_matrix.nzmax);

	std::swap(sparse_matrix.i, other.sparse_matrix.i);
	std::swap(sparse_matrix.p, other.sparse_matrix.p);
	std::swap(sparse_matrix.x, other.sparse_matrix.x);
}

// Dtor
CSparseMatrix::~CSparseMatrix() { internal_free_mem(); }
/** Erase all previous contents and leave the matrix as a "triplet" 1x1 matrix
 * without any data. */
void CSparseMatrix::clear(size_t nRows, size_t nCols)
{
	// Free old data:
	internal_free_mem();

	// Init as 1x1 triplet:
	sparse_matrix.nzmax = 1;
	sparse_matrix.m = nRows;
	sparse_matrix.n = nCols;
	sparse_matrix.i = (int*)malloc(sizeof(int) * sparse_matrix.nzmax);
	sparse_matrix.p = (int*)malloc(sizeof(int) * (sparse_matrix.n + 1));
	sparse_matrix.x = (double*)malloc(sizeof(double) * sparse_matrix.nzmax);
	sparse_matrix.nz = 0;  // >=0 -> triplet format.
}

/** free buffers (deallocate the memory of the i,p,x buffers) */
void CSparseMatrix::internal_free_mem()
{
	cs_free(sparse_matrix.i);
	cs_free(sparse_matrix.p);
	cs_free(sparse_matrix.x);
}

/** Initialization from a triplet "cs", which is first compressed */
void CSparseMatrix::construct_from_triplet(const cs& triplet)
{
	cs* sm = cs_compress(&triplet);
	copy_fast(sm);
	cs_spfree(sm);	// This will release just the "cs" structure itself, not the
	// internal buffers, by now set to NULL.
}

// To be called by constructors only, assume previous pointers are trash and
// overwrite them:
// This ONLY allocate the memory, doesn't fill it.
void CSparseMatrix::construct_from_existing_cs(const cs& sm)
{
	ASSERTMSG_(
		sm.nz == -1,
		"I expected a column-compressed sparse matrix, not a triplet form.");
	sparse_matrix.i = (int*)malloc(sizeof(int) * sm.nzmax);
	sparse_matrix.p = (int*)malloc(sizeof(int) * (sm.n + 1));
	sparse_matrix.x = (double*)malloc(sizeof(double) * sm.nzmax);
}

/** Create an initially empty sparse matrix, in the "triplet" form.
 *  Notice that you must call "compressFromTriplet" after populating the matrix
 * and before using the math operatons on this matrix.
 *  The initial size can be later on extended with insert_entry() or
 * setRowCount() & setColCount().
 */
CSparseMatrix::CSparseMatrix(size_t nRows, size_t nCols)
{
	sparse_matrix.nzmax = 1;
	sparse_matrix.m = nRows;
	sparse_matrix.n = nCols;
	sparse_matrix.i = (int*)malloc(sizeof(int) * sparse_matrix.nzmax);
	sparse_matrix.p = (int*)malloc(sizeof(int) * (sparse_matrix.n + 1));
	sparse_matrix.x = (double*)malloc(sizeof(double) * sparse_matrix.nzmax);
	sparse_matrix.nz = 0;
}

/** Insert an element into a "cs", return false on error. */
void CSparseMatrix::insert_entry(size_t row, size_t col, const double val)
{
	if (!isTriplet())
		THROW_EXCEPTION(
			"insert_entry() is only available for sparse matrix in 'triplet' "
			"format.");
	if (!cs_entry(&sparse_matrix, row, col, val))
		THROW_EXCEPTION(
			"Error inserting element in sparse matrix (out of mem?)");
}

/** Copy operator from another existing object */
CSparseMatrix& CSparseMatrix::operator=(const CSparseMatrix& other)
{
	if (&other == this) return *this;

	cs_free(sparse_matrix.i);
	cs_free(sparse_matrix.p);
	cs_free(sparse_matrix.x);

	sparse_matrix.i = (int*)malloc(sizeof(int) * other.sparse_matrix.nzmax);
	sparse_matrix.p = (int*)malloc(sizeof(int) * (other.sparse_matrix.n + 1));
	sparse_matrix.x =
		(double*)malloc(sizeof(double) * other.sparse_matrix.nzmax);
	copy(&other.sparse_matrix);
	return *this;
}

void CSparseMatrix::add_AB(const CSparseMatrix& A, const CSparseMatrix& B)
{
	ASSERT_(A.cols() == B.cols() && A.rows() == B.rows());

	cs* sm = cs_add(&(A.sparse_matrix), &(B.sparse_matrix), 1, 1);
	ASSERT_(sm);
	this->copy_fast(sm);
	cs_spfree(sm);
}

void CSparseMatrix::multiply_AB(const CSparseMatrix& A, const CSparseMatrix& B)
{
	ASSERT_(A.cols() == B.rows());

	cs* sm = cs_multiply(&(A.sparse_matrix), &(B.sparse_matrix));
	ASSERT_(sm);
	this->copy_fast(sm);
	cs_spfree(sm);
}

void CSparseMatrix::matProductOf_Ab(
	const mrpt::math::CVectorDouble& b,
	mrpt::math::CVectorDouble& out_res) const
{
	ASSERT_EQUAL_(int(b.size()), int(cols()));
	out_res.resize(rows());
	const double* y = &(b[0]);
	double* x = &(out_res[0]);
	cs_gaxpy(&sparse_matrix, y, x);
}

CSparseMatrix CSparseMatrix::transpose() const
{
	cs* sm = cs_transpose(&sparse_matrix, 1);
	ASSERT_(sm);
	CSparseMatrix SM(sm);
	cs_spfree(sm);
	return SM;
}

/** Static method to convert a "cs" structure into a dense representation of the
 * sparse matrix.
 */
void CSparseMatrix::cs2dense(const cs& SM, CMatrixDouble& d_M)
{
	d_M.setZero(SM.m, SM.n);
	if (SM.nz >= 0)	 // isTriplet ??
	{  // It's in triplet form.
		for (int idx = 0; idx < SM.nz; ++idx)
			d_M(SM.i[idx], SM.p[idx]) +=
				SM.x[idx];	// += since the convention is that duplicate (i,j)
		// entries add to each other.
	}
	else
	{  // Column compressed format:
		ASSERT_(SM.x);	// JL: Could it be nullptr and be OK???

		for (int j = 0; j < SM.n; j++)
		{
			const int p0 = SM.p[j];
			const int p1 = SM.p[j + 1];
			for (int p = p0; p < p1; p++)
				d_M(SM.i[p], j) += SM.x[p];	 // += since the convention is that
			// duplicate (i,j) entries add to
			// each other.
		}
	}
}

void CSparseMatrix::get_dense(CMatrixDouble& d_M) const
{
	cs2dense(sparse_matrix, d_M);
}

void CSparseMatrix::compressFromTriplet()
{
	if (!isTriplet())
		THROW_EXCEPTION(
			"compressFromTriplet(): Matrix is already in column-compressed "
			"format.");

	cs* sm = cs_compress(&this->sparse_matrix);
	copy_fast(sm);
	cs_spfree(sm);	// This will release just the "cs" structure itself, not the
	// internal buffers, now set to NULL.
}

/** save as a dense matrix to a text file \return False on any error.
 */
bool CSparseMatrix::saveToTextFile_dense(const std::string& filName)
{
	CMatrixDouble dense;
	this->get_dense(dense);
	try
	{
		dense.saveToTextFile(filName);
		return true;
	}
	catch (...)
	{
		return false;
	}
}

// Format:
//   NUM_ROWS   NUM_COLS   NUM_NON_ZERO_MAX
//   row_1   col_1   value_1
//   row_2   col_2   value_2
bool CSparseMatrix::saveToTextFile_sparse(const std::string& filName)
{
	FILE* f = fopen(filName.c_str(), "wt");
	if (!f) return false;

	// Help notes:
	fprintf(
		f,
		"\
# This sparse matrix can be loaded in Octave/Matlab as follows:\n\
# D=load('file.txt');\n\
# SM=spconvert(D(2:end,:));\n\
#  or...\n\
# m=D(1,1); n=D(1,2); nzmax=D(1,3);\n\
# Di=D(2:end,1); Dj=D(2:end,2); Ds=D(2:end,3);\n\
# SM=sparse(Di,Dj,Ds, m,n, nzmax);\n\n");

	// First line:
	fprintf(
		f, "%i %i %i\n", sparse_matrix.m, sparse_matrix.n,
		sparse_matrix.nzmax);  // Rows, cols, nzmax

	// Data lines:
	if (sparse_matrix.nz >= 0)	// isTriplet ??
	{  // It's in triplet form.
		for (int i = 0; i < sparse_matrix.nzmax; i++)
			if (sparse_matrix.x[i] != 0)
				fprintf(
					f, "%4i %4i %e\n", 1 + sparse_matrix.i[i],
					1 + sparse_matrix.p[i], sparse_matrix.x[i]);
	}
	else
	{  // Column compressed format:
		ASSERT_(sparse_matrix.x);  // JL: Could it be nullptr and be OK???

		for (int j = 0; j < sparse_matrix.n; j++)
		{
			const int p0 = sparse_matrix.p[j];
			const int p1 = sparse_matrix.p[j + 1];
			for (int p = p0; p < p1; p++)
				fprintf(
					f, "%4i %4i %e\n", 1 + sparse_matrix.i[p], 1 + j,
					sparse_matrix.x[p]);
		}
	}

	fclose(f);
	return true;
}

// ===============  START OF:   CSparseMatrix::CholeskyDecomp  inner class
// ==============================

/** Constructor from a square semidefinite-positive sparse matrix.
 *   The actual Cholesky decomposition takes places in this constructor.
 *  \exception std::runtime_error On non-square input matrix.
 *  \exception mrpt::math::CExceptionNotDefPos On non-semidefinite-positive
 * matrix as input.
 */
CSparseMatrix::CholeskyDecomp::CholeskyDecomp(const CSparseMatrix& SM)
	: m_symbolic_structure(nullptr),
	  m_numeric_structure(nullptr),
	  m_originalSM(&SM)
{
	ASSERT_(SM.cols() == SM.rows());
	ASSERT_(SM.isColumnCompressed());

	// symbolic decomposition:
	m_symbolic_structure =
		cs_schol(1 /* order */, &m_originalSM->sparse_matrix);

	// numeric decomposition:
	m_numeric_structure =
		cs_chol(&m_originalSM->sparse_matrix, m_symbolic_structure);

	if (!m_numeric_structure)
		throw mrpt::math::CExceptionNotDefPos(
			"CSparseMatrix::CholeskyDecomp: Not positive definite matrix.");
}

// Destructor:
CSparseMatrix::CholeskyDecomp::~CholeskyDecomp()
{
	cs_nfree(m_numeric_structure);
	cs_sfree(m_symbolic_structure);
}

/** Return the L matrix (L*L' = M), as a dense matrix. */
void CSparseMatrix::CholeskyDecomp::get_L(CMatrixDouble& L) const
{
	CSparseMatrix::cs2dense(*m_numeric_structure->L, L);
}

/** Return the vector from a back-substitution step that solves: Ux=b   */
void CSparseMatrix::CholeskyDecomp::backsub(
	const CVectorDouble& b, CVectorDouble& sol) const
{
	ASSERT_(b.size() > 0);
	sol.resize(b.size());
	this->backsub(&b[0], &sol[0], b.size());
}

/** Return the vector from a back-substitution step that solves: Ux=b   */
void CSparseMatrix::CholeskyDecomp::backsub(
	const double* b, double* sol, size_t N) const
{
	ASSERT_(N > 0);
	std::vector<double> tmp(N);

	cs_ipvec(
		m_symbolic_structure->pinv, &b[0], &tmp[0], N); /* tmp = PERMUT*b */
	// permute con. pivoting
	cs_lsolve(m_numeric_structure->L, &tmp[0]); /* tmp = L\tmp */
	cs_ltsolve(m_numeric_structure->L, &tmp[0]); /* tmp = L'\tmp */
	cs_pvec(
		m_symbolic_structure->pinv, &tmp[0], &sol[0],
		N); /* sol = PERMUT'*tmp */
	// unpermute con. pivoting

	// Result is now in "sol".
}

/** Update the Cholesky factorization from an updated vesion of the original
 * input, square definite-positive sparse matrix.
 *  NOTE: This new matrix MUST HAVE exactly the same sparse structure than the
 * original one.
 */
void CSparseMatrix::CholeskyDecomp::update(const CSparseMatrix& new_SM)
{
	ASSERTMSG_(
		m_originalSM->sparse_matrix.nzmax == new_SM.sparse_matrix.nzmax,
		"New matrix doesn't have the same sparse structure!");
	ASSERTMSG_(
		m_originalSM->sparse_matrix.n == new_SM.sparse_matrix.n,
		"New matrix doesn't have the same sparse structure!");

	m_originalSM = &new_SM;	 // Just copy the reference.

	// Release old data:
	cs_nfree(m_numeric_structure);
	m_numeric_structure = nullptr;

	// numeric decomposition:
	m_numeric_structure =
		cs_chol(&m_originalSM->sparse_matrix, m_symbolic_structure);
	if (!m_numeric_structure)
		throw mrpt::math::CExceptionNotDefPos(
			"CholeskyDecomp::update: Not positive definite matrix.");
}

// ===============    END OF: CSparseMatrix::CholeskyDecomp  inner class
// ==============================
