#include <math.h>
#include <conio.h>
#include <assert.h>
#include <Libra.h> // C++ , C

//
// Do A * B matrix multiplication, where A and B each represent a batch of matrices.
//
// Further, A is a sparse matrix, where matrices are packed along the diagonal not to interfere with each other.
//

void matrix_multiplication_batch_processing();

int main(int argc, char* argv[])
{
	libra_Init(argc, argv);
	libra_SetDefaultDataType(GFLOAT32);

	// See example on how to batch small matrices into bigger batches for fast processing.
	matrix_multiplication_batch_processing();

	libra_Shutdown();

	printf("\n\n\nPress a key...\n\n\n");
	_getch();

	return 0;
}

void getDataCompletedCallback(float* memory, void* userData)
{
	//
	//
	// ...
}

void matrix_multiplication_batch_processing()
{
	// Batch matrices into a supermatrix (sparse).
	const int matrixRowComponentCount = 4;
	const int matrixColumnComponentCount = 4;
	const int matrixCount = 10000;
	const int batchMatrixRowComponentCount = matrixCount*matrixRowComponentCount;
	const int batchMatrixColumnComponentCount = matrixCount*matrixColumnComponentCount;
	const int batchVectorRowComponentCount = batchMatrixRowComponentCount;
	const int batchVectorColumnComponentCount = matrixColumnComponentCount;
		
    //
	// Setup dummy application format
	//

	struct Matrix
	{
		float data[matrixRowComponentCount*matrixColumnComponentCount];
	};

	Matrix* applicationMatrices = new Matrix[matrixCount];

	float one = 1.0f;
	int initValue = *(int*)(void*)&one;
	// Create dummy app data
	for (int i=0;i<matrixCount;i++)
	{
		Matrix& matrix = applicationMatrices[i];
		
		for (unsigned int row=0;row<matrixRowComponentCount;++row)
		{
			for (unsigned int column = 0;column<matrixColumnComponentCount;++column)
			{
				matrix.data[row*matrixColumnComponentCount+column] = 1.0f;
			}
		}
	}

	float* values = (float*)libra_malloc(matrixRowComponentCount*matrixColumnComponentCount*matrixCount*sizeof(float));
	float* currentValue = values;
	int* rows = (int*)libra_malloc(matrixRowComponentCount*matrixColumnComponentCount*matrixCount*sizeof(int));
	int* columns = (int*)libra_malloc(matrixRowComponentCount*matrixColumnComponentCount*matrixCount*sizeof(int));

	int nz = 0;
	unsigned int rowOffset = 0;
	unsigned int columnOffset = 0;

	//
	// Copy from application format to columnmajor float array to construct a sparse matrix containing all matrices.
	//
	// See also Libra documentation.
	//
	for (unsigned int mat=0;mat<matrixCount;++mat)
	{
		for (unsigned int row=0;row<matrixRowComponentCount;++row)
		{
			for (unsigned int column = 0;column<matrixColumnComponentCount;++column)
			{
				rows[nz] = row+rowOffset;
				columns[nz] = column+columnOffset;
				*currentValue = applicationMatrices[mat].data[column];
			
				++currentValue;
				++nz;
			 }
		}
		rowOffset += matrixRowComponentCount;
		columnOffset += matrixColumnComponentCount;
	}

	// Create sparse matrix passing in row indices, column indices and values (application matrix data).
	gVar A = sparse(values, rows, columns, nz);

	// Create the right operand matrix B containing application matrices to multiply with A.
	gVar B = gVar(batchVectorRowComponentCount, batchVectorColumnComponentCount);

	printf("Sparse [%d * %d] x [%d * %d] float elements\nOr %d %dx%d matrices (jobs/objects/transforms).\n\nMatrix batches A and B size : = %f mb\n", A.rowCount(), A.columnCount(), B.rowCount(), B.columnCount(), A.rowCount()/matrixRowComponentCount, matrixRowComponentCount, matrixColumnComponentCount, ((A.rowCount()*matrixRowComponentCount*matrixColumnComponentCount*sizeof(float))/(1024.0f*1024.0f))*2.0f);
	printf("\nMatrix data elements = 1.0f\n");
	// Initialize B with application data.
	B.setData(&applicationMatrices[0].data[0]);
	// Warmup.
	gVar C = A*B;A*B;

	// Run simulation for 20 iterations.
	unsigned int iterationCount = 20;

	double startTime = libra_GetTime();

	for (unsigned int j=0;j<iterationCount;++j)
	{                               
		// Execute batched matrix multiplications
		C = A * B;                  

		// Download answer 
		// Download transfer time is hidden by computations by supplying a callback function to getData().
		C.getData(&applicationMatrices[0].data[0], 0, 0, getDataCompletedCallback, 0);
	}

	double endTime = libra_GetTime();

	printf("\nLibra average iteration time (computations + data transfer): %f ms\n", (endTime - startTime)*1000.0/(double)(iterationCount));   
	printf("\nResult matrix dim = [%d x %d], sum(C) = %f\n", C.rowCount(), C.columnCount(), (double)sum(C));   

	libra_free(values);
	libra_free(rows);
	libra_free(columns);
	delete[] applicationMatrices;
}