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

// Макросы для индексации элементов матрицы.
#define MatrixA(i,j) MatrixA[i*Size+j]
#define MatrixB(i,j) MatrixB[i*Size+j]
#define MatrixC(i,j) MatrixC[i*Size+j]
#define MatrixAA(i,j) MatrixAA[i*Size+j]
#define MatrixCC(i,j) MatrixCC[i*Size+j]

void MatrixInitialization(double*, double*, double*, int);
int VerificationResult(double*, double*, double*, int);

void main(int argc, char** argv)
{
MPI_Init(&argc,&argv);
int n,rank,Size;
MPI_Comm_rank(MPI_COMM_WORLD,&rank);
MPI_Comm_size(MPI_COMM_WORLD,&n);

Size = 512; // Размер матрицы по умолчанию.

// Присвоение значения размера матрицы из параметра командной строки.
if(argc == 2) Size = atoi(argv[1]);
                   
if(rank == 0) printf("Size = %d.\n", Size);

int i,j,k;
unsigned long int StartTime, EndTime;
double *MatrixA,*MatrixB,*MatrixC,*MatrixAA,*MatrixCC;

// Выделение памяти для хранения матриц.
MatrixB = new double [Size*Size];
if(rank == 0)
{
    MatrixA = new double [Size*Size];
    MatrixC = new double [Size*Size];
    MatrixAA = new double [Size*Size/n];
    MatrixCC = new double [Size*Size/n];
}
else
{
    MatrixAA = new double [Size*Size/n];
    MatrixCC = new double [Size*Size/n];
}

// Инициализация матриц.
if(rank == 0) MatrixInitialization(MatrixA, MatrixB, MatrixC, Size);

// Широковещательная рассылка матрицы B.
MPI_Bcast(MatrixB,Size*Size,MPI_DOUBLE,0,MPI_COMM_WORLD);

// Равномерное распределение строк матрицы А между процессами группы MPI_COMM_WORLD.
MPI_Scatter(MatrixA,Size*Size/n,MPI_DOUBLE,MatrixAA,Size*Size/n,MPI_DOUBLE,0,MPI_COMM_WORLD);

if(rank == 0) StartTime = GetTickCount();

// Перемножение Size/n строк матрицы А на матрицу B.
for (i=0; i<Size/n; i++)
{
    for (j=0; j<Size; j++)
    {
    MatrixCC(i,j)=0;
    for (k=0; k<Size; k++)
        MatrixCC(i,j) = MatrixCC(i,j) + MatrixAA(i,k)*MatrixB(k,j);
    }
}

MPI_Barrier(MPI_COMM_WORLD);

if(rank == 0) EndTime = GetTickCount();

// Сборка результатов перемножения со всех процессов группы MPI_COMM_WORLD.
MPI_Gather(MatrixCC,Size*Size/n,MPI_DOUBLE,MatrixC,Size*Size/n,MPI_DOUBLE,0,MPI_COMM_WORLD);
if(rank == 0)
{
    printf("Time multiplying: %.2f sec.\n",(float)(EndTime-StartTime)/1000);
    if(VerificationResult(MatrixA, MatrixB, MatrixC, Size)) printf("Data is correct.\n");
    else printf("Data is incorrect.\n");
}

delete[] MatrixB;
if(rank == 0)
{
    delete[] MatrixA;
    delete[] MatrixC;
    delete[] MatrixAA;
    delete[] MatrixCC;
}
else
{
    delete[] MatrixAA;
    delete[] MatrixCC;
}

MPI_Finalize();
}

void MatrixInitialization(double* MatrixA, double* MatrixB, double* MatrixC, int Size)
{
int i,j;
for (i=0; i<Size; i++)
{
    for (j=0; j<Size; j++)
    {
    MatrixA(i,j) = 1;
    MatrixB(i,j) = 2;
    MatrixC(i,j) = 0;
    }
}
}

int VerificationResult(double* MatrixA, double* MatrixB, double* MatrixC, int Size)
{
int i,j,k;
double** MatrixEtalon = new double* [Size];
for (i=0; i<Size; i++) MatrixEtalon[i] = new double [Size];

for (i=0; i<Size; i++)
{
    for (j=0; j<Size; j++)
    {
    MatrixEtalon[i][j]=0;
    for (k=0; k<Size; k++)
        MatrixEtalon[i][j] = MatrixEtalon[i][j] + MatrixA(i,k)*MatrixB(k,j);
        if(MatrixEtalon[i][j]!=MatrixC(i,j)) return 0;
    }
}

for (i=0; i<Size; i++) delete[] MatrixEtalon[i];
delete[] MatrixEtalon;

return 1;
}