/* squic_interface.hpp
 *
 * Copyright (C) Matthias Bollhoefer
 *
 * 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 <https://www.gnu.org/licenses/>.
 */

#include "squic.hpp"
/**
 * @brief      Creat buffer at given location
 *
 * @param      buffer  The buffer
 * @param[in]  length  The length of the buffer
 */
__attribute__ ((visibility("default"))) void SQUIC_CPP_UTIL_memset_integer(integer *&buffer, integer length) {
	buffer = new integer[length]();
}

/**
 * @brief      Creat buffer at given location
 *
 * @param      buffer  The buffer
 * @param[in]  length  The length of the buffer
 */
__attribute__ ((visibility("default"))) void SQUIC_CPP_UTIL_memset_double(double *&buffer, integer length) {
	buffer = new double[length]();
}


/**
 * @brief      Creat a new copy of the buffer
 *
 * @param      buffer  The copy buffer
 * @param      values  The buffer to be copied
 * @param[in]  length  The length of the buffer
 */
__attribute__ ((visibility("default"))) void SQUIC_CPP_UTIL_memcopy_integer(integer *&buffer, integer *values, integer length) {
	buffer = new integer[length]();

	for (integer i = 0; i < length; ++i) {
		buffer[i] = values[i];
	}
}

/**
 * @brief      Creat a new copy of the buffer
 *
 * @param      buffer  The copy buffer
 * @param      values  The buffer to be copied
 * @param[in]  length  The length of the buffer
 */
__attribute__ ((visibility("default"))) void SQUIC_CPP_UTIL_memcopy_double(double *&buffer, double *values, integer length) {
	buffer = new double[length]();

	for (integer i = 0; i < length; ++i) {
		buffer[i] = values[i];
	}
}

/**
 * @brief      Free the buffer
 *
 * @param      buffer  The buffer
 */
__attribute__ ((visibility("default"))) void SQUIC_CPP_UTIL_memfree_integer(integer *&buffer) {
	delete [] buffer;
}

/**
 * @brief      Free the buffer
 *
 * @param      buffer  The buffer
 */
__attribute__ ((visibility("default"))) void SQUIC_CPP_UTIL_memfree_double(double *&buffer) {
	delete [] buffer;
}

/**
 * @brief SQUIC C++ Interface
 *
 * This function implements the SQUIC algorithm in C++.
 *
 * @param[in]  mode          Runtime mode values [0,1,2,3,4] for the "block" algorithm and [5,6,7,8,9] for the "scalar" algorithm. 
 *                           (Note: Mode "0" is the ONLY supported value).
 * @param[in]  p             Number of random variables.
 * @param[in]  n             Number of samples.
 * @param[in]  Y             Pointer to the data.
 * @param[in]  lambda        Scalar tuning parameter.
 * @param[in]  M_rinx        Row indices of M matrix.
 * @param[in]  M_cptr        Column pointers of M matrix.
 * @param[in]  M_val         Values of M matrix.
 * @param[in]  M_nnz         Number of non-zero elements in M matrix.
 * @param[in]  max_iter      Maximum number of Newton iterations. 
 *                           (Note: Setting max_iter to 0 returns the sample covariance matrix in output iC).
 * @param[in]  inv_tol       Inversion tolerance for approximate inversion. 
 *                           (Note: inv_tol should be positive and it is recommended that inv_tol is equal to term_tol).
 * @param[in]  term_tol      Termination tolerance.
 *                           (Note: term_tol should be positive and it is recommended that inv_tol is equal to term_tol).
 * @param[in]  verbose       Verbose level (0 or 1).
 * @param[in/out]  X_rinx    Row indices of the precision matrix. 
 *                           (Note: Initial value X0_rinx should be passed here).
 * @param[in/out]  X_cptr    Column pointers of the precision matrix.
 *                           (Note: Initial value X0_cptr should be passed here).
 * @param[in/out]  X_val     Values of the precision matrix.
 *                           (Note: Initial value X0_val should be passed here).
 * @param[in/out]  X_nnz     Number of non-zero elements in the precision matrix.
 *                           (Note: Initial value X_nnz should be passed here).
 * @param[in/out]  W_rinx    Row indices of the covariance matrix.
 *                           (Note: Initial value W0_rinx should be passed here).
 * @param[in/out]  W_cptr    Column pointers of the covariance matrix.
 *                           (Note: Initial value W0_rinx should be passed here).
 * @param[in/out]  W_val     Values of the covariance matrix.
 *                           (Note: Initial value W0_rinx should be passed here).
 * @param[in/out]  W_nnz     Number of non-zero elements in the covariance matrix.
 *                           (Note: Initial value W0_rinx should be passed here).
 * @param[out] info_num_iter Number of Newton iterations performed.
 * @param[out] info_times    Six-element array containing computation times for: 
 *                           1) Total, 2) Sample covariance, 3) Optimization, 4) Factorization, 5) Approximate inversion, 6) Coordinate update.
 * @param[out] info_objective Array containing the objective function values.
 *                            (Note: This array must be of length max_iter when passed in. Upon output, only info_num_iter elements will be written).
 * @param[out] info_logdetX  Log determinant of X.
 * @param[out] info_trSX     Trace of SX.
 */
__attribute__ ((visibility("default"))) void SQUIC_CPP(
    int mode,
    integer p, integer n, double *Y,
    double lambda,
    integer *M_rinx, integer *M_cptr, double *M_val, integer M_nnz,
    int max_iter, double inv_tol, double term_tol, int verbose,
    integer *&X_rinx, integer *&X_cptr, double *&X_val, integer &X_nnz,
    integer *&W_rinx, integer *&W_cptr, double *&W_val, integer &W_nnz,
    int &info_num_iter,
    double *&info_times,	 //length must be 6: [time_total,time_impcov,time_optimz,time_factor,time_aprinv,time_updte]
    double *&info_objective, // length must be size max_iter
    double &info_logdetX,
    double &info_trSX) {

	assert(p > 2);
	assert(n > 1);


	assert(X_nnz >= p);
	assert(W_nnz >= p);

	assert(lambda   > 0.0 );
	assert(max_iter >= 0  );
	assert(inv_tol  >= 0.0);
	assert(term_tol >= 0.0);
	assert(mode     >= 0 && mode <= 9);

	verbose = min(2, max(0, verbose));


	// kind of (incomplete) Cholesky decomposition: 0 (CHOLMOD), 1 (PARDISO), 2 (JANUS)
	int cholesky = 0;
	int coord_update;
	int inversion;

	bool use_ZERO_M = M_nnz < 1;
	bool use_block = mode < 5;


	double *mu = new double[p];
	double *sd = new double[p];


	// block variant
	SQUIC *sq = new SQUIC(n, p, Y);
	sq->RunTimeConfig.verbose = verbose;
	sq->RunTimeConfig.generate_scm = DETERMINISTIC; 
	sq->RunTimeConfig.block_upd = COLLECTIVE_CDU;
        sq->RunTimeConfig.factorization=CHOL_SUITESPARSE;
        sq->RunTimeConfig.inversion=BLOCK_FSPAI;

	
	// Using scalar lambda and I intial guess


	if (use_ZERO_M) {
		SQUIC::SparseMatrix X0;
		SQUIC::SparseMatrix W0;
		sq->CSC_to_CustomeData(X_rinx, X_cptr, X_val, X_nnz, &X0);
		sq->CSC_to_CustomeData(W_rinx, W_cptr, W_val, W_nnz, &W0);
		sq->run(lambda, inv_tol, max_iter, term_tol, &X0, &W0);
		sq->FreeSparse(&X0);
		sq->FreeSparse(&W0);
	} else {
		SQUIC::SparseMatrix X0;
		SQUIC::SparseMatrix W0;
		SQUIC::SparseMatrix LambdaMatrix;
		sq->CSC_to_CustomeData(X_rinx, X_cptr, X_val, X_nnz, &X0);
		sq->CSC_to_CustomeData(W_rinx, W_cptr, W_val, W_nnz, &W0);
		sq->CSC_to_CustomeData(M_rinx, M_cptr, M_val, M_nnz, &LambdaMatrix);
		sq->run(lambda, &LambdaMatrix, inv_tol, max_iter, term_tol, &X0, &W0);
		sq->FreeSparse(&X0);
		sq->FreeSparse(&W0);
		sq->FreeSparse(&LambdaMatrix);
	}

	//Re allocate the memory
	delete[] X_rinx;
	delete[] X_cptr;
	delete[] X_val;
	delete[] W_rinx;
	delete[] W_cptr;
	delete[] W_val;

	if (max_iter == 0) {
		X_nnz = sq->X.nnz;
		W_nnz = sq->S.nnz; // nnz for S!!!

		X_rinx = new integer[X_nnz];
		X_cptr = new integer[p + 1];
		X_val = new double[X_nnz];

		W_rinx = new integer[W_nnz];
		W_cptr = new integer[p + 1];
		W_val = new double[W_nnz];

		X_cptr[0] = 0;
		W_cptr[0] = 0;

		integer nnz_count_X = 0;
		integer nnz_count_W = 0;
		double val;

		for (integer j = 0; j < p; j++) {

			// Copy value of X
			for (integer k = 0; k < sq->X.ncol[j]; k++) {
				val = sq->X.val[j][k];
				if (fabs(val) > EPS) {
					X_rinx[nnz_count_X] = sq->X.rowind[j][k];
					X_val[nnz_count_X] = val;
					nnz_count_X++;
				}
			}
			X_cptr[j + 1] = nnz_count_X;

			// Copy values of S to W!!!
			for (integer k = 0; k < sq->S.ncol[j]; k++) {
				val = sq->S.val[j][k];
				if (fabs(val) > EPS) {
					W_rinx[nnz_count_W] = sq->S.rowind[j][k];
					W_val[nnz_count_W] = val;
					nnz_count_W++;
				}
			}
			W_cptr[j + 1] = nnz_count_W;
		}

		// The memory footprint is the same but the nnz values in the front of the array
		X_nnz = nnz_count_X;
		W_nnz = nnz_count_W; // nnz for S!!!

		// Passout out stats
		info_num_iter = sq->Stat.time_itr.size();
		info_times[0] = sq->Stat.time_total;
		info_times[1] = sq->Stat.time_cov;
		info_times[2] = 0.0;
		info_times[3] = 0.0;
		info_times[4] = 0.0;
		info_times[5] = 0.0;

		// we do nothing with info_objective[]
		//for (integer i = 0; i < sq->Stat.time_itr.size(); ++i)
		//{
		//	info_objective[i] = sq->Stat.opt[i];
		//}

		// info_logdetX = do nothing ;
		// info_trSX = do nothing ;
	} else {
		// General case: we passout both X and W
		X_nnz = sq->X.nnz;
		W_nnz = sq->W.nnz;

		X_rinx = new integer[X_nnz];
		X_cptr = new integer[p + 1];
		X_val = new double[X_nnz];

		W_rinx = new integer[W_nnz];
		W_cptr = new integer[p + 1];
		W_val = new double[W_nnz];

		X_cptr[0] = 0;
		W_cptr[0] = 0;

		integer nnz_count_X = 0;
		integer nnz_count_W = 0;
		double val;

		for (integer j = 0; j < p; j++) {


			// For column j we go
			for (integer k = 0; k < sq->X.ncol[j]; k++) {
				val = sq->X.val[j][k];
				if (fabs(val) > EPS) {
					X_rinx[nnz_count_X] = sq->X.rowind[j][k];
					X_val[nnz_count_X] = val;
					nnz_count_X++;
				}
			}
			X_cptr[j + 1] = nnz_count_X;

			// For column j we go
			for (integer k = 0; k < sq->W.ncol[j]; k++) {
				val = sq->W.val[j][k];
				if (fabs(val) > EPS) {
					W_rinx[nnz_count_W] = sq->W.rowind[j][k];
					W_val[nnz_count_W] = val;
					nnz_count_W++;
				}
			}
			W_cptr[j + 1] = nnz_count_W;
		}

		// The memory footprint is the same but the nnz values in the front of the array.
		X_nnz = nnz_count_X;
		W_nnz = nnz_count_W;

		// Passout out stats
		info_num_iter = sq->Stat.opt.size(); // Number of objective values evalutions is the number of iterations
		info_times[0] = sq->Stat.time_total;
		info_times[1] = sq->Stat.time_cov;

		info_times[2] = accumulate(sq->Stat.time_itr.begin(), sq->Stat.time_itr.end(), 0.0);
		info_times[3] = accumulate(sq->Stat.time_chol.begin(), sq->Stat.time_chol.end(), 0.0);
		info_times[4] = accumulate(sq->Stat.time_inv.begin(), sq->Stat.time_inv.end(), 0.0);
		info_times[5] = accumulate(sq->Stat.time_upd.begin(), sq->Stat.time_upd.end(), 0.0);

		for (integer i = 0; i < info_num_iter; ++i) {
			info_objective[i] = sq->Stat.opt[i];
		}
		info_logdetX = sq->Stat.logdetX;
		info_trSX = sq->Stat.trSX;
	}

	sq->SQUIC_DeleteVars();
	sq->FreeSparse(&sq->X);
	sq->FreeSparse(&sq->W);

	delete sq;
};
