#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>


#define MASTER				0
#define MATRIX_RANK			200		/* column and row size of the matrix */
#define MATRIX_MPI_TYPE		MPI_FLOAT
typedef int MATRIX_DATATYPE;

void initializeMatrix(MATRIX_DATATYPE m1[MATRIX_RANK][MATRIX_RANK], \
	MATRIX_DATATYPE m2[MATRIX_RANK][MATRIX_RANK]);

void printMatrix(MATRIX_DATATYPE m1[MATRIX_RANK][MATRIX_RANK]);
void multiplyMatrix(int offset, int chunkSize, MATRIX_DATATYPE *m1, \
	MATRIX_DATATYPE m2[MATRIX_RANK][MATRIX_RANK], MATRIX_DATATYPE *mr);

int main(int argc, char **argv)
{
	MATRIX_DATATYPE matrixA[MATRIX_RANK][MATRIX_RANK] = {0}; 
	MATRIX_DATATYPE matrixB[MATRIX_RANK][MATRIX_RANK] = {0};
	MATRIX_DATATYPE matrixResult[MATRIX_RANK][MATRIX_RANK] = {0};
	int numTasks, taskId, offset, chunkSize, source, dest;
	int tag1, tag2, tag3, i;
	MPI_Status status;

	MPI_Init(&argc, &argv);
	MPI_Comm_size(MPI_COMM_WORLD, &numTasks);
	if (numTasks > MATRIX_RANK ||  MATRIX_RANK % numTasks != 0) {
		printf("Quitting. Number of MPI tasks must devide %d.\n", MATRIX_RANK);
	    MPI_Abort(MPI_COMM_WORLD, 1);
		exit(1);
	}
	MPI_Comm_rank(MPI_COMM_WORLD, &taskId);
	tag1 = 1; // tag for sending offset
	tag2 = 2; // tag for sending part of matrix A
	tag3 = 3; // tag fpr sending part of matrix B
	chunkSize = MATRIX_RANK / numTasks;

	/* master task */
	if (taskId == MASTER) {
		/* initialise matrix data */
		initializeMatrix(matrixA, matrixB);
		offset = chunkSize;
		/* send data and offset to the rest of tasks */
		for (dest = 1; dest < numTasks; dest++) {
			MPI_Send(&offset, 1, MPI_INT, dest, tag1, MPI_COMM_WORLD);
			MPI_Send(&matrixA[offset][0], chunkSize*MATRIX_RANK, MATRIX_MPI_TYPE, dest, tag2, MPI_COMM_WORLD);
			MPI_Send(matrixB, MATRIX_RANK*MATRIX_RANK, MATRIX_MPI_TYPE, dest, tag3, MPI_COMM_WORLD);
			offset += chunkSize;
		}

		/* Master task does its own work */
		offset = 0;
		multiplyMatrix(offset, chunkSize, &matrixA[offset][0], matrixB, &matrixResult[offset][0]);

		/* Wait to receive results from each task */
		for(source = 1; source < numTasks; source++) {
			MPI_Recv(&offset, 1, MPI_INT, source, tag1, MPI_COMM_WORLD, &status);
			MPI_Recv(&matrixResult[offset][0], chunkSize*MATRIX_RANK, MATRIX_MPI_TYPE, source, tag2, MPI_COMM_WORLD, &status);
		}
		/* print the result */
		//printMatrix(matrixResult);
	}

	/* non-master tasks */
	if (taskId > MASTER) {
		source = MASTER;
		MPI_Recv(&offset, 1, MPI_INT, source, tag1, MPI_COMM_WORLD, &status);
		MPI_Recv(&matrixA[offset][0], chunkSize*MATRIX_RANK, MATRIX_MPI_TYPE, source, tag2, MPI_COMM_WORLD, &status);
		MPI_Recv(matrixB, MATRIX_RANK*MATRIX_RANK, MATRIX_MPI_TYPE, source, tag3, MPI_COMM_WORLD, &status);
#if 0
		printf("task %d received, offset = %d\n", taskId, offset);
		for (i = 0; i < MATRIX_RANK; i++) {
			printf("%f  ", matrixA[offset][i]);
		}
		printf("\n");
#endif 		
		/* does calculation */
		multiplyMatrix(offset, chunkSize, &matrixA[offset][0], matrixB, &matrixResult[offset][0]);
		/* send the result back to master */
		MPI_Send(&offset, 1, MPI_INT, MASTER, tag1, MPI_COMM_WORLD);
		MPI_Send(&matrixResult[offset][0], chunkSize*MATRIX_RANK, MATRIX_MPI_TYPE, MASTER, tag2, MPI_COMM_WORLD);
	}

	MPI_Finalize();
	return 0;
}

void multiplyMatrix(int offset, int chunkSize, MATRIX_DATATYPE *m1, \
	MATRIX_DATATYPE m2[MATRIX_RANK][MATRIX_RANK], MATRIX_DATATYPE *mr)
{
	int i, row, col, sum = 0;
	int total = MATRIX_RANK * chunkSize;

	for (i = 0; i < total; i++) {
		sum = 0;
		row = (int) i / MATRIX_RANK; 
		for (col = 0; col < MATRIX_RANK; col++) {
			sum += *(m1+row*MATRIX_RANK+col) * m2[col][i%MATRIX_RANK];
		}
		*(mr+i) = sum;
	}
}

void initializeMatrix(MATRIX_DATATYPE m1[MATRIX_RANK][MATRIX_RANK], \
	MATRIX_DATATYPE m2[MATRIX_RANK][MATRIX_RANK])
{
	int i, j;

	for (i=0; i < MATRIX_RANK; i++) {
		for (j=0; j < MATRIX_RANK; j++) {
			m1[j][i] = m2[j][i] = i+j;
		}
	}
}

void printMatrix(MATRIX_DATATYPE m1[MATRIX_RANK][MATRIX_RANK])
{

	int i, j;

	for (i=0; i < MATRIX_RANK; i++) {
		for (j=0; j < MATRIX_RANK; j++) {
			printf("%f  ", m1[j][i]);
		}
		printf("\n");
	}
}