#include "stdio.h"
#include <iostream>
#include <mpi.h>

using namespace std;

int size = 1;
int rank = 0;

double * Matrix (int Rows, int Columns)
{
	double * matrix;
	try
	{
		matrix = new double [Rows * Columns];
	}
	catch (int err_code)
	{
		cout << "Memory allocation error " << err_code << endl;
		delete [] matrix;
		matrix = NULL;
		return NULL;
	}
	return matrix;
}

void DeleteMatrix(double * matrix)
{
	delete [] matrix;
	matrix = NULL;
}

void FillMatrix(double * matrix, int Rows, int Columns, int InitValue)
{
	for (int i = 0; i < Rows; i++)
		for (int j = 0; j < Columns; j++)
			matrix[i * Columns + j] = i + j + InitValue;
	return;
}

void PrintMatrix(double * matrix, int Rows, int Columns)
{
	int Rows_1 = Rows - 1;
	int Columns_1 = Columns - 1;
	cout << "[";
	for (int i = 0; i < Rows; i++)
	{
		cout << "[";
		for (int j = 0; j < Columns_1; j++)
		{
			cout << matrix[i * Columns + j] << ", ";
		}
		cout << matrix [i * Columns + Columns_1] << "]";
		if (i != Rows_1)
			cout << ", ";
	}
	cout << "]" << endl;
	return;
}

//	note* that Rows2 = Columns1 and that's why it is omitted
void Multiply(double * A, double * B, int Rows1, int Columns1, int Columns2,  double * Result)
{
	int N;
	for (int i = 0; i < Rows1; i++)
		for (int j = 0; j < Columns2; j++)
		{
			N = i * Columns2 + j;
			Result[N] = 0;
			for (int h = 0; h < Columns1; h++)
				Result[N] += A[i * Columns1 + h] * B[h * Columns2 + j];
		}
	return;
}


void MultiplyParallel(double * A, double * B, int Rows1, int Columns1, int Columns2, double * Result, int root)
{
	int rows1 = Rows1 / size;
	int columns2 = Columns2 / size;
	int columns1 = Columns1;
	int rows2 = Columns1;
	int offset;
	int next_process;
	int prev_process;
	
	//	special type for transmitting columns
	int * blocklengths = new int [rows2];
	int * displacements = new int [rows2];
	MPI_Datatype * types = new MPI_Datatype [rows2];

	for (int k = 0; k < rows2; k++)
	{
		blocklengths[k] = columns2;
		displacements[k] = sizeof(double) * Columns2 * k;
		types[k] = MPI_DOUBLE;
	}

	MPI_Datatype special_type, uspecial_type;
	MPI_Type_create_struct(rows2, blocklengths, displacements, types, &uspecial_type);
	MPI_Type_commit(&uspecial_type);
	MPI_Type_create_resized(uspecial_type, 0, sizeof(double) * columns2, &special_type);
	MPI_Type_commit(&special_type);
	//	end of special type

	double * A1 = Matrix(rows1, columns1);
	double * B1 = Matrix(rows2, columns2);
	double * Result1 = Matrix(rows1, Columns2);
	double * Result2 = Matrix(rows1, columns2);

	//	initial data -> all processes
	MPI_Scatter(A, rows1 * columns1, MPI_DOUBLE, A1, rows1 * columns1, MPI_DOUBLE, root, MPI_COMM_WORLD);
	MPI_Scatter(B, 1, special_type, B1, rows2 * columns2, MPI_DOUBLE, root, MPI_COMM_WORLD); 

	//	computing and exchanging columns
	for (int k = 0; k < size; k++)
	{
		offset = ((rank + k) * columns2)%Columns2;
		double ntime=MPI_Wtime();
		cout << rank << " s " ;
		printf("%7f \n",ntime);
		ntime=MPI_Wtime();
		Multiply(A1, B1, rows1, columns1, columns2, Result2);
		cout << rank << " e " ;
		printf("%7f \n",ntime);
		for (int i = 0; i < rows1; i++)
			for (int j = 0; j < columns2; j++)
				Result1[i * Columns2 + j + offset] = Result2[i * columns2 + j];
		next_process = (rank + 1)%size;
		prev_process = (rank - 1);
		if (prev_process < 0) prev_process += size;
		MPI_Sendrecv_replace(B1, rows2 * columns2, MPI_DOUBLE, prev_process, 0, next_process, MPI_ANY_TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
	}
	//	results from all processes -> root process
	MPI_Gather(Result1, rows1 * Columns2, MPI_DOUBLE, Result, rows1 * Columns2, MPI_DOUBLE, root, MPI_COMM_WORLD);

	DeleteMatrix(Result2);
	DeleteMatrix(Result1);
	DeleteMatrix(A1);
	DeleteMatrix(B1);
	MPI_Type_free(&special_type);
	MPI_Type_free(&uspecial_type);
	return;
}

int main(int argc, char ** argv)
{
	MPI_Init(&argc, &argv);
	MPI_Comm_size(MPI_COMM_WORLD, &size);
	MPI_Comm_rank(MPI_COMM_WORLD, &rank);
	int msize = 1024;
	double time = 0;

	if (rank == 0) 
	{
		double * A = Matrix(msize, msize);
		double * B = Matrix(msize, msize);
		double * C = Matrix(msize, msize);
		FillMatrix(A, msize, msize, 0);
		FillMatrix(B, msize, msize, 1);
		//PrintMatrix(A, msize, msize);
		//PrintMatrix(B, msize, msize);
		time = MPI_Wtime();
		MultiplyParallel(A, B, msize, msize, msize, C, 0);
		
		time= MPI_Wtime() - time;
		cout << "Result of parallel computation: (time is " << time << ")" << endl;
		//PrintMatrix(C, msize, msize);

		//	check ;)
		time = MPI_Wtime();
		Multiply(A, B, msize, msize, msize, C);
		time = MPI_Wtime() - time;
		cout << "Really result must be: (time is " << time << ")" << endl;
		//PrintMatrix(C, msize, msize);

		DeleteMatrix(A);
		DeleteMatrix(B);
		DeleteMatrix(C);
	}
	else MultiplyParallel(NULL, NULL, msize, msize, msize, NULL, 0);

	MPI_Finalize();
	return 0;
}

