#include "util.h"

template<typename ComplexOrRealType>
void csc_kron_mult_method(const int imethod,
                          const int nrow_A,
                          const int ncol_A,
                          const PsimagLite::Vector<int>::Type& acolptr,
                          const PsimagLite::Vector<int>::Type& arow,
                          const typename PsimagLite::Vector<ComplexOrRealType>::Type& aval,
                          const int nrow_B,
                          const int ncol_B,
                          const PsimagLite::Vector<int>::Type& bcolptr,
                          const PsimagLite::Vector<int>::Type& brow,
                          const typename PsimagLite::Vector<ComplexOrRealType>::Type& bval,
                          const PsimagLite::Matrix<ComplexOrRealType>& yin,
                          PsimagLite::Matrix<ComplexOrRealType>& xout)
{
	const int nrow_X = nrow_B;
	const int ncol_X = nrow_A;

	const int nrow_Y = ncol_B;
	const int ncol_Y = ncol_A;

	int nnz_A = csc_nnz( ncol_A, acolptr );
	int nnz_B = csc_nnz( ncol_B, bcolptr );
	int has_work = (nnz_A >= 1) && (nnz_B >= 1);

	assert((imethod == 1) ||
	       (imethod == 2) ||
	       (imethod == 3));

	if (!has_work) {
		return;
	};
	/*
 *   -------------------------------------------------------------
 *   A and B are in compressed sparse COLUMN format
 *
 *   X += kron( A, B) * Y
 *   that can be computed as either
 *   imethod == 1
 *
 *   X(ib,ia) += (B(ib,jb) * Y(jb,ja) ) * transpose(A(ia,ja)   or
 *               BY(ib,ja) = B(ib,jb)*Y(jb,ja)
 *               BY is nrow_B by ncol_A, need   2*nnz(B)*ncolA flops
 *
 *   X(ib,ia) +=   BY(ib,ja) * transpose(A(ia,ja)) need 2*nnz(A)*nrowB flops
 *
 *   imethod == 2
 *
 *   X(ib,ia) += B(ib,jb) * (Y(jb,ja) * transpose(A))    or
 *                YAt(jb,ia) = Y(jb,ja) * transpose(A(ia,ja))
 *                YAt is ncolB by nrowA, need 2*nnz(A) * ncolB flops
 *
 *   X(ib,ia) += B(ib,jb) * YAt(jb,ia)  need nnz(B) * nrowA flops
 *
 *   imethod == 3
 *
 *   X += kron(A,B) * Y   by visiting all non-zero entries in A, B
 *
 *   this is feasible only if A and B are very sparse, need nnz(A)*nnz(B) flops
 *   -------------------------------------------------------------
 */



	if (imethod == 1) {

	/*
	 *  --------------------------------------------
	 *  BY(ib,ja) = B(ib,jb)*Y(jb,ja)
	 *
	 *  X(ib,ia) += BY(ib,ja ) * transpose(A(ia,ja))
	 *  --------------------------------------------
	 */



	/*
	 * ---------------
	 * setup BY(ib,ja)
	 * ---------------
	 */
		const int nrow_BY = nrow_B;
		const int ncol_BY = ncol_A;
		PsimagLite::Matrix<ComplexOrRealType> by_(nrow_BY, ncol_BY);

		{
			int iby = 0;
			int jby = 0;

			// not needed, FIXME
			for( jby=0; jby < ncol_BY; jby++) {
				for( iby=0; iby < nrow_BY; iby++) {
					by_(iby,jby) = 0;
				};
			};
		}


		{
			/*
	 * ------------------------------
	 * BY(ib,ja)  = B(ib,jb)*Y(jb,ja)
	 * ------------------------------
	 */
			const char trans = 'N';
			csc_matmul_pre(  trans,
			                 nrow_B,
			                 ncol_B,
			                 bcolptr,
			                 brow,
			                 bval,

			                 nrow_Y,
			                 ncol_Y,
			                 yin,

			                 nrow_BY,
			                 ncol_BY,
			                 by_);

		}

		{
			/*
	 * -------------------------------------------
	 * X(ib,ia) += BY(ib,ja) * transpose(A(ia,ja))
	 * -------------------------------------------
	 */
			const char trans = 'T';

			csc_matmul_post(
			            trans,
			            nrow_A,
			            ncol_A,
			            acolptr,
			            arow,
			            aval,

			            nrow_BY,
			            ncol_BY,
			            by_,

			            nrow_X,
			            ncol_X,
			            xout);

		}
	}
	else if (imethod == 2) {
		/*
	 * ---------------------
	 * YAt(jb,ia) = Y(jb,ja) * tranpose(A(ia,ja))
	 * X(ib,ia) += B(ib,jb) * YAt(jb,ia)
	 * ---------------------
	 */



		/*
	 * ----------------
	 * setup YAt(jb,ia)
	 * ----------------
	 */

		PsimagLite::Matrix<ComplexOrRealType> yat_(ncol_B, nrow_A);

		int nrow_YAt = ncol_B;
		int ncol_YAt = nrow_A;


		{
			int jb = 0;
			int ia = 0;

			//not needed, FIXME
			for(ia=0; ia < ncol_A; ia++) {
				for(jb=0; jb < ncol_B; jb++) {
					yat_(jb,ia) = 0;
				};
			};
		}


		{
			/*
	  * ---------------------
	  * YAt(jb,ia) = Y(jb,ja) * tranpose(A(ia,ja)
	  * ---------------------
	  */
			const char transa = 'T';



			csc_matmul_post( transa,
			                 nrow_A,
			                 ncol_A,
			                 acolptr,
			                 arow,
			                 aval,

			                 nrow_Y,
			                 ncol_Y,
			                 yin,

			                 nrow_YAt,
			                 ncol_YAt,
			                 yat_);
		}




		{
			/*
	 * ------------
	 * X(ib,ia) += B(ib,jb) * YAt(jb,ia)
	 * ------------
	 */

			const char trans = 'N';
			csc_matmul_pre( trans,
			                nrow_B,
			                ncol_B,
			                bcolptr,
			                brow,
			                bval,

			                nrow_YAt,
			                ncol_YAt,
			                yat_,

			                nrow_X,
			                ncol_X,
			                xout);

		}
	}
	else if (imethod == 3) {
	/*
	* ---------------------------------------------
	* C = kron(A,B)
	* C([ib,ia], [jb,ja]) = A(ia,ja)*B(ib,jb)
	* X([ib,ia]) += C([ib,ia],[jb,ja]) * Y([jb,ja])
	* ---------------------------------------------
	*/

		int ja = 0;
		for(ja=0; ja < ncol_A; ja++) {
			int istarta = acolptr[ja];
			int ienda = acolptr[ja+1]-1;
			int ka = 0;
			for(ka=istarta; ka <= ienda; ka++) {
				int ia = arow[ka];
				ComplexOrRealType aij = aval[ka];
				int jb = 0;
				for(jb=0; jb < ncol_B; jb++) {
					int istartb = bcolptr[jb];
					int iendb = bcolptr[jb+1]-1;
					int kb = 0;
					for(kb=istartb; kb <= iendb; kb++) {
						int ib = brow[kb];
						ComplexOrRealType bij = brow[kb];
						ComplexOrRealType cij = aij * bij;

						xout(ib,ia) +=  (cij * yin(jb,ja));
					};
				};
			};
		};

	};

}

template<typename ComplexOrRealType>
void csc_kron_mult( const int nrow_A,
                    const int ncol_A,
                    const PsimagLite::Vector<int>::Type& acolptr,
                    const PsimagLite::Vector<int>::Type& arow,
                    const typename PsimagLite::Vector<ComplexOrRealType>::Type& aval,
                    const int nrow_B,
                    const int ncol_B,
                    const PsimagLite::Vector<int>::Type& bcolptr,
                    const PsimagLite::Vector<int>::Type& brow,
                    const typename PsimagLite::Vector<ComplexOrRealType>::Type& bval,
                    const PsimagLite::Matrix<ComplexOrRealType>& yin,
                    PsimagLite::Matrix<ComplexOrRealType>& xout,
                    const typename PsimagLite::Real<ComplexOrRealType>::Type denseFlopDiscount)

{
/*
 *   -------------------------------------------------------------
 *   A and B are in compressed sparse COLUMN format
 *
 *   X += kron( A, B) * Y
 *   that can be computed as either
 *   imethod == 1
 *
 *   X(ib,ia) += (B(ib,jb) * Y(jb,ja) ) * transpose(A(ia,ja)   or
 *               BY(ib,ja) = B(ib,jb)*Y(jb,ja)
 *               BY is nrow_B by ncol_A, need   2*nnz(B)*ncolA flops
 *
 *   X(ib,ia) +=   BY(ib,ja) * transpose(A(ia,ja)) need 2*nnz(A)*nrowB flops
 *
 *   imethod == 2
 *
 *   X(ib,ia) += B(ib,jb) * (Y(jb,ja) * transpose(A))    or
 *                YAt(jb,ia) = Y(jb,ja) * transpose(A(ia,ja))
 *                YAt is ncolB by nrowA, need 2*nnz(A) * ncolB flops
 *
 *   X(ib,ia) += B(ib,jb) * YAt(jb,ia)  need nnz(B) * nrowA flops
 *
 *   imethod == 3
 *
 *   X += kron(A,B) * Y   by visiting all non-zero entries in A, B
 *
 *   this is feasible only if A and B are very sparse, need nnz(A)*nnz(B) flops
 *   -------------------------------------------------------------
 */

	int nnz_A = csc_nnz( ncol_A, acolptr );
	int nnz_B = csc_nnz( ncol_B, bcolptr );
	int has_work = (nnz_A >= 1) && (nnz_B >= 1);

	ComplexOrRealType kron_nnz = 0;
	ComplexOrRealType kron_flops = 0;
	int imethod = 1;

	if (!has_work) {
		return;
	};

	estimate_kron_cost( nrow_A,ncol_A,nnz_A, nrow_B,ncol_B,nnz_B,
	                    &kron_nnz, &kron_flops, &imethod, denseFlopDiscount);



	csc_kron_mult_method(
	            imethod,
	            nrow_A,
	            ncol_A,
	            acolptr,
	            arow,
	            aval,

	            nrow_B,
	            ncol_B,
	            bcolptr,
	            brow,
	            bval,

	            yin,
	            xout );
}
#undef BY
#undef YAt
#undef X
#undef Y
