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

std::ostream& operator <<(std::ostream &os, MPI::Status& status);

void allocateMatrix(double**& pr_matrix, const unsigned& pr_rowCount, const 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()
{
	const int MASTER_RANK = 0;
	double **matrixA, **matrixB, **matrixC;
	unsigned matrixARowCount, matrixAColumnCount, matrixBRowCount, matrixBColumnCount, matrixCRowCount, matrixCColumnCount;
	double *matrixBColumn = 0;

	// inicijalizuj MPI
	MPI::Init();
	// odredi svoj poredak unutar MPI sveta (prvi je 0)
	int rank = MPI::COMM_WORLD.Get_rank();
	// odredi veličinu MPI sveta
	int 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 ili dimenzije ne odgovaraju broju procesa, kraj sveta
		std::cout << "Ja sam " << rank << " od " << size << ", vidim sledece vrednosti: A " << matrixARowCount << "X" << matrixAColumnCount << ", B" << matrixBRowCount << "X" << matrixBColumnCount << std::endl;
		if (matrixAColumnCount != matrixBRowCount || size != matrixBColumnCount)
		{
			std::cerr << "Kraj sveta\n";
			MPI::COMM_WORLD.Abort(15);
		}
		else
		{
			allocateMatrix(matrixC, matrixARowCount, matrixBRowCount);
			writeMatrix(rank, matrixC, matrixARowCount, matrixBRowCount);
		}
	}

	std::cout << "Ja sam " << rank << " od " << size << ", vidim sledece vrednosti: A " << matrixARowCount << "X" << matrixAColumnCount << ", B" << matrixBRowCount << "X" << matrixBColumnCount << std::endl;

	// ako je izvršavanje stiglo do ovde, svi treba da dobiju 
	// 1. dimenzije matrica A i B
	MPI::COMM_WORLD.Bcast(&matrixAColumnCount, 1, MPI::UNSIGNED, MASTER_RANK);
	MPI::COMM_WORLD.Bcast(&matrixARowCount, 1, MPI::UNSIGNED, MASTER_RANK);
	MPI::COMM_WORLD.Bcast(&matrixBColumnCount, 1, MPI::UNSIGNED, MASTER_RANK);
	MPI::COMM_WORLD.Bcast(&matrixBRowCount, 1, MPI::UNSIGNED, MASTER_RANK);

	std::cout << "Ja sam " << rank << " od " << size << ", vidim sledece vrednosti: A " << matrixARowCount << "X" << matrixAColumnCount << ", B" << matrixBRowCount << "X" << matrixBColumnCount << std::endl;

	// 2. elemente matrice A
	// oni koji nisu "master" prvo moraju da alociraju odgovarajuće bafere
	if (MASTER_RANK != rank)
	{
		allocateMatrix(matrixA, matrixARowCount, matrixAColumnCount);
	}

	std::cout << "Ja sam " << rank << " od " << size << ", vidim sledece vrednosti: A " << matrixARowCount << "X" << matrixAColumnCount << ", B" << matrixBRowCount << "X" << matrixBColumnCount << std::endl;
	writeMatrix(rank, matrixA, matrixARowCount, matrixAColumnCount);

	// pošto je matrica dinamički alocirana, 
	// redovi najverovatnije nisu poređani jedan za drugim, zato šalji red po red
	for (unsigned i = 0; i < matrixARowCount; i++)
	{
		MPI::COMM_WORLD.Bcast(matrixA[i], matrixAColumnCount, MPI::DOUBLE, MASTER_RANK);
	}

	std::cout << "Ja sam " << rank << " od " << size << ", vidim sledece vrednosti: A " << matrixARowCount << "X" << matrixAColumnCount << ", B" << matrixBRowCount << "X" << matrixBColumnCount << std::endl;
	writeMatrix(rank, matrixA, matrixARowCount, matrixAColumnCount);

	// 3. odgovarajuću kolonu matrice B
	// svi moraju da se pripreme za prijem svoje kolone matrice B
	matrixBColumn = new double[matrixBRowCount];
	// master će slati red po red, tako da svaki proces dobije elemente iz svoje kolone
	for (unsigned i = 0; i < matrixBRowCount; i++)
	{
		MPI::COMM_WORLD.Scatter(matrixB[i], 1, MPI::DOUBLE, &matrixBColumn[i], 1, MPI::DOUBLE, MASTER_RANK);
	}

	std::cout << "Ja sam " << rank << " od " << size << ", vidim sledece vrednosti: A " << matrixARowCount << "X" << matrixAColumnCount << ", B" << matrixBRowCount << "X" << matrixBColumnCount << std::endl;
	writeArray(rank, matrixBColumn, matrixBRowCount);

	// svako množi redove iz matrice A sa svojom kolonom
	double *tempMatrixCColumn = new double[matrixARowCount];
	for (unsigned i = 0; i < matrixARowCount; i++)
	{
		tempMatrixCColumn[i] = multiplyRowAndColumn(matrixA[i], matrixBColumn, matrixBRowCount);
	}
	writeArray(rank, tempMatrixCColumn, matrixARowCount);

	// dimenzije matrice C određene su dimenzijama matrica A i B
	matrixCRowCount = matrixARowCount;
	matrixCColumnCount = matrixBColumnCount;

	// potrebno je sakupiti sve rezultate u "master"
	for (unsigned i = 0; i < matrixCRowCount; i++)
	{
		MPI::COMM_WORLD.Gather(&tempMatrixCColumn[i], 1, MPI::DOUBLE, matrixC[i], 1, MPI::DOUBLE, MASTER_RANK);
	}

	// na kraju, "master" ispisuje dobijenu matricu
	if (MASTER_RANK == rank)
	{
		writeMatrix(rank, matrixC, matrixCRowCount, matrixCColumnCount);
	}

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

	return 0;
}

std::ostream& operator <<(std::ostream &os, MPI::Status& status)
{
	/// @todo treba promeniti ovaj operator u dumpStatus funkciju,
	/// ciji ce izlaz biti zavisan od vrste podatka koji su primljeni
	os << "posiljalac " << status.Get_source()
		<< ", oznaka " << status.Get_tag()
		<< ", kod greske " << status.Get_error() 
		<< ", duzina poruke " << status.Get_count(MPI::INT) << " podataka tipa MPI::INT";
	return os;
}