/************************************************************************/
/* Ovo je samo jedno od mnogih mogućih ispravnih rešenja                */
/************************************************************************/
#include <mpi.h>
#include <iostream>
#include <iomanip>

enum Tags 
{ 
	ROW_TAG = 1000, 
	COLUMN_TAG,
	RESULT_TAG,
	ID_TAG
};

void allocateMatrix(double**& pr_matrix, unsigned& pr_rowCount, unsigned& pr_columnCount)
{
	pr_matrix = new double*[pr_rowCount];
	for (unsigned i = 0; i < pr_rowCount; i++)
	{
		pr_matrix[i] = new double[pr_columnCount];
	}
}

void loadAndAllocateMatrix(double**& pr_matrix, unsigned& pr_rowCount, unsigned& pr_columnCount)
{
	std::cout << "Unesite dimenzije matrice: ";
	std::cin >> pr_rowCount >> pr_columnCount;
	allocateMatrix(pr_matrix, pr_rowCount, pr_columnCount);

	std::cout << "Unesite elemente matrice:\n";
	for (unsigned i = 0; i < pr_rowCount; i++)
	{
		for (unsigned j = 0; j < pr_columnCount; j++)
		{
			std::cin >> pr_matrix[i][j];
		}
	}
}

void writeMatrix(int rank, double**& pr_matrix, unsigned& pr_rowCount, unsigned& pr_columnCount)
{
	std::cout << "Elementi matrice u procesu " << rank << std::endl;
	for (unsigned i = 0; i < pr_rowCount; i++)
	{
		for (unsigned j = 0; j < pr_columnCount; j++)
		{
			std::cout << std::setw(5) << pr_matrix[i][j];
		}
		std::cout << std::endl;
	}
}

void writeArray(int rank, double *pp_array, unsigned pv_elementCount)
{
	std::cout << "Elementi niza u procesu " << rank << std::endl;
	for (unsigned i = 0; i < pv_elementCount; i++)
	{
		std::cout << std::setw(5) << pp_array[i];
	}
	std::cout << std::endl;
}

void destroyMatrix(double**& pr_matrix, unsigned pr_rowCount)
{
	for (unsigned i = 0; i < pr_rowCount; i++)
	{
		delete [] pr_matrix[i], pr_matrix[i] = 0;
	}
	delete [] pr_matrix, pr_matrix = 0;
}

double multiplyRowAndColumn(double* pp_row, double* pp_column, unsigned pv_length)
{
	double result = 0;
	for (unsigned i = 0; i < pv_length; i++)
	{
		result += pp_row[i] * pp_column[i];
	}
	return result;
}

int main(int argc, char *argv[])
{
	double **matrixA = 0, **matrixB = 0, **matrixC = 0, *matrixARow = 0, *matrixBColumn = 0;
	unsigned matrixARowCount, matrixAColumnCount, matrixBRowCount, matrixBColumnCount, matrixCRowCount, matrixCColumnCount, matrixCElementCount;
	const int MASTER_RANK = 0;

	// inicijalizuj MPI
	MPI::Init(argc, argv);
	// odredi svoj poredak unutar MPI sveta (prvi je 0)
	unsigned rank = MPI::COMM_WORLD.Get_rank();
	// odredi veličinu MPI sveta
	unsigned size = MPI::COMM_WORLD.Get_size();

	// ako si "master", učitaj matrice i pošalji svima potrebne podatke
	if (MASTER_RANK == rank)
	{
		loadAndAllocateMatrix(matrixA, matrixARowCount, matrixAColumnCount);
		loadAndAllocateMatrix(matrixB, matrixBRowCount, matrixBColumnCount);
		// ako se matrice ne poklapaju, kraj sveta
		if (matrixAColumnCount != matrixBRowCount || size < matrixARowCount*matrixBColumnCount + 2)
		{
			std::cerr << "Kraj sveta\n";
			MPI::COMM_WORLD.Abort(2806);
		}
		else
		{
			allocateMatrix(matrixC, matrixARowCount, matrixBRowCount);
		}
	}

	// ako je izvršavanje stiglo do ovde, svi treba da 
	// 1. dobiju neophodne dimenzije matrica i odrede ostale na osnovu njih
	MPI::COMM_WORLD.Bcast(&matrixARowCount, 1, MPI::UNSIGNED, 0);
	MPI::COMM_WORLD.Bcast(&matrixAColumnCount, 1, MPI::UNSIGNED, 0);
	matrixBRowCount = matrixAColumnCount;
	MPI::COMM_WORLD.Bcast(&matrixBColumnCount, 1, MPI::UNSIGNED, 0);
	matrixCRowCount = matrixARowCount;
	matrixCColumnCount = matrixBColumnCount;
	matrixCElementCount = matrixCRowCount * matrixCColumnCount;

	// 2. radnici (oni koji će obavljati množenje) prvo moraju da alociraju odgovarajuće bafere
	if (MASTER_RANK != rank && rank <= matrixCElementCount)
	{
		// potrebno je alocirati samo jedan red matrice A u procesima radnicima
		matrixARow = new double[matrixAColumnCount];
		// isto važi i za matricu B - potrebna je samo jedna kolona po procesu radniku
		matrixBColumn = new double[matrixBRowCount];
	}

	// 3. dobiju potrebne podatke
	if (MASTER_RANK == rank)
	{
		for (unsigned rowIdx = 0; rowIdx < matrixARowCount; rowIdx++)
			for (unsigned colIdx = 0; colIdx < matrixBColumnCount; colIdx++)
			{
				int destinationRank = 1 + rowIdx * matrixCColumnCount + colIdx;
				MPI::COMM_WORLD.Send(matrixA[rowIdx], matrixAColumnCount, MPI::DOUBLE, destinationRank, ROW_TAG);
				double *matrixBColumn = new double[matrixBRowCount];
				for (unsigned tempIdx = 0; tempIdx < matrixBRowCount; tempIdx++)
				{
					matrixBColumn[tempIdx] = matrixB[tempIdx][colIdx];
				}
				MPI::COMM_WORLD.Send(matrixBColumn, matrixBRowCount, MPI::DOUBLE, destinationRank, COLUMN_TAG);
				delete [] matrixBColumn;
			}
	} 
	else if (MASTER_RANK != rank && rank <= matrixCElementCount)
	{
		MPI::COMM_WORLD.Recv(matrixARow, matrixAColumnCount, MPI::DOUBLE, MASTER_RANK, ROW_TAG);
		MPI::COMM_WORLD.Recv(matrixBColumn, matrixBRowCount, MPI::DOUBLE, MASTER_RANK, COLUMN_TAG);
	}

	// 4. dobiju odgovarajući novi komunikator, boja odgovara indeksu vrste
	int color;
	if (MASTER_RANK != rank && rank <= matrixCElementCount)
		color = (rank-1) / matrixCColumnCount;
	else
		color = matrixCRowCount; // master i suvišni procesi ostaju u istom novom komunikatoru

	// Standard: "If all the keys are specified in the same way, then all the
	// processes in a given color will have the relative rank order as they 
	// did in their parent group. (In general, they will have different ranks.)"
	// Zato će procesi završiti tačno kako je potrebno narednom programskom kodu
	int key = 0;
	MPI::Intracomm newComm = MPI::COMM_WORLD.Split(color, key);

	// svi koji treba da rade računanje počinju sa tim
	if (color != matrixCRowCount)
	{
		double *matrixCRow = new double[matrixCColumnCount];
		double matrixCElement = multiplyRowAndColumn(matrixARow, matrixBColumn, matrixAColumnCount);

		std::cout << "Ja sam " << rank << " (boja " << color << ", rang " << newComm.Get_rank() << "), moj element je " << matrixCElement << std::endl;

		// sakupi sve rezultate iz tekućeg reda i pošalji masteru uz svoj rank u COMM_WORLD
		newComm.Gather(&matrixCElement, 1, MPI::DOUBLE, matrixCRow, 1, MPI::DOUBLE, MASTER_RANK);
		if (newComm.Get_rank() == MASTER_RANK)
		{
			std::cout << "Ja sam lokalni master " << rank << " (boja " << color << ", rang " << newComm.Get_rank() << "), moja vrsta je..." << std::endl;
			writeArray(rank, matrixCRow, matrixCColumnCount);
			MPI::COMM_WORLD.Send(&color, 1, MPI::INT, MASTER_RANK, ID_TAG);
			MPI::COMM_WORLD.Send(matrixCRow, matrixCColumnCount, MPI::DOUBLE, MASTER_RANK, RESULT_TAG);
		}
		delete [] matrixARow;
		delete [] matrixBColumn;
		delete [] matrixCRow;
	}
	// master sakuplja jedan po jedan rezultat, ostali ne rade ništa
	else if (MASTER_RANK == rank)
	{
		for (unsigned rowIdx = 0; rowIdx < matrixCRowCount; rowIdx++)
		{
			int currentRow;
			MPI::Status status;
			MPI::COMM_WORLD.Recv(&currentRow, 1, MPI::INT, MPI::ANY_SOURCE, ID_TAG, status);
			MPI::COMM_WORLD.Recv(matrixC[currentRow], matrixCColumnCount, MPI::DOUBLE, status.Get_source(), RESULT_TAG);
		}
		// prikaži rezultat
		writeMatrix(MASTER_RANK, matrixC, matrixCRowCount, matrixCColumnCount);
		destroyMatrix(matrixA, matrixARowCount);
		destroyMatrix(matrixB, matrixBRowCount);
		destroyMatrix(matrixC, matrixCRowCount);
	}
	else
	{
		std::cout << "Ja sam " << rank << ". Nerad me ubi, molim Vas da naredni put zadate mnogo posla" << std::endl;
	}

	newComm.Free();

	// odjava iz MPI sveta
	MPI::Finalize();	

	return 0;
}