/*===========================================================================*
* kronprod.c
* Fast mixed Kronecker product on real or complex matrices.
*
* Copyright 2009 Levente Hunyadi
* All rights reserved.
*===========================================================================*/

#define __PACKAGENAME__ "math:kronprod"
#include "common.c"

/* Two-dimensional indexing in arrays. */
#define el(M,i,j,m,n) M[j*m+i]

/** Mixed-Kronecker product of real matrices A, B and C.
*  The result is equivalent to computing kron(A,B)*C but is significantly faster.
*  @param[in] A A matrix with dimensions a_m * a_n.
*  @param[in] B A matrix with dimensions b_m * b_n.
*  @param[in] C A sparse matrix with dimensions (a_n*b_n) * c_n.
*  @return A matrix with dimensions (a_m*b_m) * c_n.
*/
mxArray* kronprod_real(const mxArray* A, const mxArray* B, const mxArray* C) {
	const double* pa = mxGetPr(A);
	mwSize a_m = mxGetM(A);
	mwSize a_n = mxGetN(A);
	mwIndex a_i, a_j;
	const double* pb = mxGetPr(B);
	mwSize b_m = mxGetM(B);
	mwSize b_n = mxGetN(B);
	mwIndex b_i, b_j;
	const double* pc = mxGetPr(C);
	mwSize c_n = mxGetN(C);
	const mwIndex* c_jc = mxGetJc(C);
	const mwIndex* c_ir = mxGetIr(C);
	mwIndex c_i;  /* row index in C */
	mwIndex c_j;  /* column index in C */
	mwIndex c_k;  /* element index in column of C */
	mxArray* D = mxCreateDoubleMatrix(a_m*b_m, c_n, mxREAL);
	double* pd = mxGetPr(D);
	double a, b, c;
	
	for (c_j = 0; c_j < c_n; c_j++) {  /* enumerate columns of sparse matrix C */
		for (a_i = 0; a_i < a_m; a_i++) {  /* enumerate rows of full matrix A */
			for (b_i = 0; b_i < b_m; b_i++) {  /* enumerate rows of full matrix B */
				double sum = 0.0;
				for (c_k = c_jc[c_j]; c_k < c_jc[c_j+1]; c_k++) {  /* enumerate elements in sparse matrix column */
					c_i = c_ir[c_k];
					a_j = c_i / b_n;  /* corresponding column index in A */
					b_j = c_i % b_n;  /* corresponding column index in B */
					a = pa[a_j*a_m + a_i];  /* full matrix indexing, A(a_i, a_j) */
					b = pb[b_j*b_m + b_i];  /* full matrix indexing, B(b_i, b_j) */
					c = pc[c_k];  /* sparse matrix indexing, C(c_i, c_i) */
					sum += a * b * c;
				}
				*pd++ = sum;  /* D(a_i*b_m, c_j) */
			}
		}
	}
	return D;
}

mxArray* kronprod_complex(const mxArray* A, const mxArray* B, const mxArray* C) {
	const double* pra = mxGetPr(A);
	const double* pia = mxGetPi(A);
	mwSize a_m = mxGetM(A);
	mwSize a_n = mxGetN(A);
	const double* prb = mxGetPr(B);
	const double* pib = mxGetPi(B);
	mwSize b_m = mxGetM(B);
	mwSize b_n = mxGetN(B);
	const double* prc = mxGetPr(C);
	const double* pic = mxGetPi(C);
	mwSize c_n = mxGetN(C);
	const mwIndex* c_jc = mxGetJc(C);
	const mwIndex* c_ir = mxGetIr(C);
	mwIndex c_i;  /* row index in C */
	mwIndex c_j;  /* column index in C */
	mwIndex c_k;  /* element index in column of C */
	mxArray* D = mxCreateDoubleMatrix(a_m*b_m, c_n, mxCOMPLEX);
	double* pd = mxGetPr(D);
	
	mexErrMsgIdAndTxt(__INVALIDVALUE__, "Operation not yet implemented.");
	return D;
}

void mexFunction(int nlhs, mxArray* plhs[], int nrhs, const mxArray* prhs[]) {
	/* check for proper number of input and output arguments */
	nargineqchk(nrhs, 3);
	nargoutltechk(nlhs, 1);

	/* check that parameters are 2d double real matrices */
	argdoublechk(prhs, 0);
	arg2dchk(prhs, 0);
	argnonemptychk(prhs, 0);
	argfullchk(prhs, 0);

	argdoublechk(prhs, 1);
	arg2dchk(prhs, 1);
	argnonemptychk(prhs, 1);
	argfullchk(prhs, 1);

	argdoublechk(prhs, 2);
	arg2dchk(prhs, 2);
	argnonemptychk(prhs, 2);
	argsparsechk(prhs, 2);

	if (mxGetN(prhs[0])*mxGetN(prhs[1]) != mxGetM(prhs[2])) {
		mexErrMsgIdAndTxt(__DIMENSIONMISMATCH__, "Inconsistent dimensions for A, B and C: size(A,2)*size(B,2) ~= size(C,1).");
	}
	
	if (!mxIsComplex(prhs[0]) && !mxIsComplex(prhs[1]) && !mxIsComplex(prhs[2])) {
		plhs[0] = kronprod_real(prhs[0], prhs[1], prhs[2]);
	} else if (mxIsComplex(prhs[0]) && mxIsComplex(prhs[1]) && mxIsComplex(prhs[2])) {
		plhs[0] = kronprod_complex(prhs[0], prhs[1], prhs[2]);
	} else {
		mexErrMsgIdAndTxt(__INVALIDVALUE__, "Operation supported only on exclusively real or exclusively complex operands.");
	}
}