#include "PsimagLite.h"
#include "util.h"

template <typename ComplexOrRealType>
void csr_submatrix(const PsimagLite::CrsMatrix<ComplexOrRealType>& a,
    const int nrow_B,
    const int ncol_B,
    const int max_nnz,
    const PsimagLite::Vector<int>::Type& rindex,
    const PsimagLite::Vector<int>::Type& cindex,
    PsimagLite::CrsMatrix<ComplexOrRealType>& b)
{
	/*
	 * ---------------------------------------------------------------------------
	 * Input:   sparse matrix A in compressed sparse row format
	 *
	 * Input:   list of row index rindex(0:(nrow_B-1)) and
	 *          list of column index cindex(0:(ncol_B-1))
	 *
	 * Out:     extract B = A( rindex(:), cindex(:)) in compress sparse row format
	 * ---------------------------------------------------------------------------
	 */

	/*
	 * ----------------------------------------------------
	 * setup boolean array for fast mapping of column index
	 * ----------------------------------------------------
	 */

	const int ncol_A = a.cols();

	int* cmap = new int[ncol_A];
	int* nnz = new int[nrow_B];

	int ja = 0;
	int ib = 0;
	int jb = 0;

	for (ja = 0; ja < ncol_A; ja++) {
		cmap[ja] = -1;
	};

	for (jb = 0; jb < ncol_B; jb++) {
		int ja = cindex[jb];
		assert((0 <= ja) && (ja < ncol_A));

		cmap[ja] = jb;
	};

	/*
	 * -------------------------------------------------------
	 * first pass to calculate number of nonzeros per row in B
	 * -------------------------------------------------------
	 */

	for (ib = 0; ib < nrow_B; ib++) {
		nnz[ib] = 0;
	};

	for (ib = 0; ib < nrow_B; ib++) {
		int ia = rindex[ib];
		int istart = a.getRowPtr(ia);
		int iend = a.getRowPtr(ia + 1);

		assert((0 <= ia) && (static_cast<SizeType>(ia) < a.rows()));

		int k = 0;
		for (k = istart; k < iend; k++) {
			int ja = a.getCol(k);
			assert((0 <= ja) && (ja < ncol_A));

			int jb = cmap[ja];
			int isok_jb = (0 <= jb) && (jb < ncol_B);
			if (isok_jb) {
				nnz[ib] += 1;
			};
		};
	};

	/*
	 * ----------------------------
	 * check for sufficient storage
	 * ----------------------------
	 */
	int total_nnz = 0;
	for (ib = 0; ib < nrow_B; ib++) {
		total_nnz += nnz[ib];
	};
	if (total_nnz > max_nnz) {
		err("Not enough storage for nnz\n");
	}

	/*
	 * ---------------------------------
	 * prefix sum to setup row pointers
	 * ---------------------------------
	 */
	const bool use_push = true;
	if (use_push) {
		b.resize(nrow_B, ncol_B);
		b.reserve(max_nnz);
	} else {
		b.resize(nrow_B, ncol_B, max_nnz);
	};

	b.setRow(0, 0);
	for (ib = 0; ib < nrow_B; ib++)
		b.setRow(ib + 1, b.getRowPtr(ib) + nnz[ib]);

	/*
	 * ------------------------
	 * reset array for 2nd pass
	 * ------------------------
	 */
	for (ib = 0; ib < nrow_B; ib++) {
		nnz[ib] = 0;
	};

	/*
	 * --------------------------------------
	 * second pass to fill in compress sparse row
	 * data structure
	 * --------------------------------------
	 */
	for (ib = 0; ib < nrow_B; ib++) {
		int ia = rindex[ib];
		int istart = a.getRowPtr(ia);
		int iend = a.getRowPtr(ia + 1);
		int k = 0;
		for (k = istart; k < iend; k++) {
			int ja = a.getCol(k);
			ComplexOrRealType aij = a.getValue(k);

			jb = cmap[ja];
			int isvalid = (0 <= jb) && (jb < ncol_B);
			if (isvalid) {
				if (use_push) {
					b.pushValue(aij);
					b.pushCol(jb);
				} else {
					int ipos = b.getRowPtr(ib) + nnz[ib];
					b.setValues(ipos, aij);
					b.setCol(ipos, jb);
					nnz[ib] += 1;
				};
			};
		};
	};

	b.checkValidity();
	delete[] cmap;
	delete[] nnz;
}
