#include <iostream>
#include <mpi.h>
#include <math.h>

int size = 1;	//	number of processes
int rank = 0;		//	rank of the current process
int current_process = 0;
int tag = 0;

using namespace std;

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;
		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 << "(" << Rows << ", " << Columns << ")[";
	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 InsertMatrixBlock(double * M, int Rows, int Columns, double * Result, int RowsR, int ColumnsR, int StartRow, int StartColumn){
	for (int i = 0; i < Rows; i++)
	{
		int _i = StartRow + i;
		for (int j = 0; j < Columns; j++)
			Result[_i * ColumnsR + j + StartColumn] = M[i * Columns + j];
	}
}

void CopyMatrixBlock(double * M, int Rows, int Columns, int StartRow, int StartColumn, double * Result, int RowsR, int ColumnsR){
	for (int i = 0; i < RowsR; i++)
	{
		int _i = StartRow + i;
		for (int j = 0; j < ColumnsR; j++)
			Result[i * ColumnsR + j] = M[_i * Columns + j + StartColumn];
	}
}

void Sum(double * A, double * B, int Rows, int Columns, double * Result)
{
	for (int i = 0; i < Rows; i++)
	{
		int _i = i * Columns;
		for (int j = 0; j < Columns; j++)
		{
			int N = _i + j;
			Result[N] = A[N] + B[N];
		}
	}
	return;
}

void Sub(double * A, double * B, int Rows, int Columns, double * Result)
{
	for (int i = 0; i < Rows; i++)
	{
		int _i = i * Columns;
		for (int j = 0; j < Columns; j++)
		{
			int N = _i + j;
			Result[N] = A[N] - B[N];
		}
	}
	return;
}

int next_process()	//	returns next computational process, for convience
{
	if (current_process < size - 1)
		current_process++;
	else
		current_process = 1;
	return current_process;
}

int next_tag()	//	returns next tag, for convience
{
	if (tag < 10000)	//	must be even
		tag++;
	else
		tag = 1;
	return tag;
}

void close_processes()
{
	bool exit = true;
	for (int num = 1; num < size; num++)
	{
		MPI_Send(&exit, 1, MPI_LOGICAL, num, 0, MPI_COMM_WORLD);
	}
}

//	only N*N matrixes !
void StrassenMultiply(double * A, double * B, int size, double * Result, const int pow = 0)	//	pow is needed for MPI... 
{
	const int half = size / 2;
	const int half2 = half * half;

	//	must be 63 
	if (half <= 63)	//	if matrix is small then use common multiple
	{
		Multiply(A, B, size, size, size, Result);
		return;
	}

	double * P [7];
	double * _A [7];
	double * _B [7];
	for (int k = 0; k < 7; k++)
	{
		P[k] = Matrix(half, half);
		_A[k] = Matrix(half, half);
		_B[k] = Matrix(half, half);
	}

	double * A11 = Matrix(half, half);
	double * A12 = Matrix(half, half);
	double * A21 = Matrix(half, half);
	double * A22 = Matrix(half, half);
	double * B11 = Matrix(half, half);
	double * B12 = Matrix(half, half);
	double * B21 = Matrix(half, half);
	double * B22 = Matrix(half, half);

	CopyMatrixBlock(A, size, size, 0, 0, A11, half, half);
	CopyMatrixBlock(A, size, size, 0, half, A12, half, half);
	CopyMatrixBlock(A, size, size, half, 0, A21, half, half);
	CopyMatrixBlock(A, size, size, half, half, A22, half, half);

	CopyMatrixBlock(B, size, size, 0, 0, B11, half, half);
	CopyMatrixBlock(B, size, size, 0, half, B12, half, half);
	CopyMatrixBlock(B, size, size, half, 0, B21, half, half);
	CopyMatrixBlock(B, size, size, half, half, B22, half, half);

	Sum(A11, A22, half, half, _A[0]);
	Sum(B11, B22, half, half, _B[0]);
	Sum(A21, A22, half, half, _A[1]);
	_B[1] = B11;
	_A[2] = A11;
	Sub(B12, B22, half, half, _B[2]);
	_A[3] = A22;
	Sub(B21, B11, half, half, _B[3]);
	Sum(A11, A12, half, half, _A[4]);
	_B[4] = B22;
	Sub(A21, A11, half, half, _A[5]);
	Sum(B11, B12, half, half, _B[5]);
	Sub(A12, A22, half, half, _A[6]);
	Sum(B21, B22, half, half, _B[6]);

	if ((rank == 0) && (size > 1))
	{
		if (pow == 1)
		{
			int ranks[7];
			int tags[14];
			for (int k = 0; k < 7; k++)
			{
				MPI_Send(_A[k], half2, MPI_DOUBLE, ranks[k] = next_process(), tags[k*2] = next_tag(), MPI_COMM_WORLD);
				MPI_Send(_B[k], half2, MPI_DOUBLE, ranks[k], tags[k*2 + 1] = next_tag(), MPI_COMM_WORLD);
			}
			for (int k = 0; k < 7; k++)
				MPI_Recv(P[k], half2, MPI_DOUBLE, ranks[k], tags[k*2], MPI_COMM_WORLD, MPI_STATUS_IGNORE);
		}
		else
		{
			int decpow = pow - 1;
			for (int k = 0; k < 7; k++)
				StrassenMultiply(_A[k], _B[k], half, P[k], decpow);
		}
	}
	else	//	other processes simply doing calculations
	{
		for (int k = 0; k < 7; k++)
			StrassenMultiply(_A[k], _B[k], half, P[k]);
	}

	double * C11 = Matrix(half, half);
	double * C12 = Matrix(half, half);
	double * C21 = Matrix(half, half);
	double * C22 = Matrix(half, half);
	Sum(P[0], P[3], half, half, C11);
	Sub(C11, P[4], half, half, C11);
	Sum(C11, P[6], half, half, C11);
	Sum(P[2], P[4], half, half, C12);
	Sum(P[1], P[3], half, half, C21);
	Sub(P[0], P[1], half, half, C22);
	Sum(C22, P[2], half, half, C22);
	Sum(C22, P[5], half, half, C22);

	InsertMatrixBlock(C11, half, half, Result, size, size, 0, 0);
	InsertMatrixBlock(C12, half, half, Result, size, size, 0, half);
	InsertMatrixBlock(C21, half, half, Result, size, size, half, 0);
	InsertMatrixBlock(C22, half, half, Result, size, size, half, half);

	for (int k = 0; k < 7; k++)
	{
		DeleteMatrix(P[k]);
		DeleteMatrix(_A[k]);
		DeleteMatrix(_B[k]);
	}

	DeleteMatrix(C11);
	DeleteMatrix(C12);
	DeleteMatrix(C21);
	DeleteMatrix(C22);

	return;
}

void comp_process()	//	called by non-root computational processes
{
	MPI_Status status;
	int _tag = 0;
	int size = 0;
	int size2 = 0;
	bool exit = false;
	while (!exit)
	{
		MPI_Probe(0, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
		size2 = status.count / sizeof (double);
		_tag = status.MPI_TAG;
		size = (int)sqrt((double)size2);
		if (size <= 1)
			exit = true;
		else
		{
			//cout << "Process " << rank << " start calculating..." << endl;
			double * A = Matrix(size, size);
			double * B = Matrix(size, size);
			double * Result = Matrix(size, size);
			MPI_Recv(A, size2, MPI_DOUBLE, 0, _tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
			MPI_Recv(B, size2, MPI_DOUBLE, 0, _tag + 1, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
			StrassenMultiply(A, B, size, Result);
			MPI_Send(Result, size2, MPI_DOUBLE, 0, _tag, MPI_COMM_WORLD);			
		}
		//cout << "Process " << rank << " end calculating." << endl;
	}
}

int main(int argc, char ** argv)
{
	MPI_Init(&argc, &argv);
	MPI_Comm_size(MPI_COMM_WORLD, &size);
	MPI_Comm_rank(MPI_COMM_WORLD, &rank);
	const int comp_size = size - 1;	//	number of computational processes (for some calculations)
	int power = 0;
	int comp = 1;
	int comp_per_process = 0;
	double time = 0;

	if (rank == 0) 
	{
		/*while (comp < comp_size) 
		{
			power++;
			comp *= 7;
		}
		if (comp_size > 0) 
		{
			comp_per_process = comp / comp_size;
			if (comp % comp_size != 0) comp_per_process++;
		}
		else 
			comp_per_process = 0;

		cout << "comp_per_process " << comp_per_process << endl;
		cout << "power " << power << endl;
		cout << "comp " << comp << endl;
		cout << "comp_size " << comp_size << endl;*/
		power = 1;

		int msize = 2048;
		double * A = Matrix(msize, msize);	//	choose only 2^n dimensions!!!
		double * B = Matrix(msize, msize);
		double * C = Matrix(msize, msize);
		double * D = Matrix(msize, msize);
		FillMatrix(A, msize, msize, 1);
		FillMatrix(B, msize, msize, 2);
		//PrintMatrix(A, msize, msize);
		//PrintMatrix(B, msize, msize);

		time = MPI_Wtime();
		StrassenMultiply(A, B, msize, C, power);
		time = MPI_Wtime() - time;
		cout << "Matrix was calculated using MPI, time is " << time << endl;
		//PrintMatrix(C, msize, msize);
		time = MPI_Wtime();
		Multiply(A, B, msize, msize, msize, D);
		time = MPI_Wtime() - time;
		cout << "Matrix was calculated directly, time is " << time << endl;
		//PrintMatrix(D, msize, msize);

		DeleteMatrix(A);
		DeleteMatrix(B);
		DeleteMatrix(C);
		DeleteMatrix(D);
		close_processes();
	}
	else
		comp_process();

	//cout << "Process #" << rank << " finished." << endl;
	MPI_Finalize();
	return 0;
}