// инициализация матриц
void InitializationMatrix(double* MatrixA, double* MatrixB, int Size)
{
 for (int i = 0; i < Size; i++)
  for (int j = 0; j < Size; j++)
  {
   MatrixA[i*Size+j] = rand()/double(100);
   MatrixB[i*Size+j] = rand()/double(100);
  }
}

void Scatter(double* Matrix, double* MatrixBlock, int Size, int BSize)
{
  double* MatrixRow = new double [BSize*Size];

  if (GridC[1] == 0)  MPI_Scatter(Matrix, BSize*Size, MPI_DOUBLE, MatrixRow, BSize*Size, MPI_DOUBLE, 0, ColComm);

  for (int i = 0; i < BSize; i++)
  MPI_Scatter(&MatrixRow[i*Size], BSize, MPI_DOUBLE, &(MatrixBlock[i*BSize]), BSize, MPI_DOUBLE, 0, RowComm);
 
  delete [] MatrixRow;
}

// разделение данных на блоки  для умножения
void Distribution(double* MatrixA, double* MatrixB, double* Block, double* BlockB, int Size, int BSize)
{
 Scatter(MatrixA, Block, Size, BSize);
 Scatter(MatrixB, BlockB, Size, BSize);
}

// создание решетки
void CreateGird()
{
 int DimSize[2]; 
 int Periodic[2];
 int Subdims[2]; 
 
 DimSize[0] = Grid;
 DimSize[1] = Grid;
 Periodic[0] = 0;
 Periodic[1] = 0;

 MPI_Cart_create(MPI_COMM_WORLD, 2, DimSize, Periodic, 1, &GridComm);

 MPI_Cart_coords(GridComm, Rank, 2, GridC);
 
 Subdims[0] = 0; 
 Subdims[1] = 1; 
 
 MPI_Cart_sub(GridComm, Subdims, &RowComm);
 
    Subdims[0] = 1;
    Subdims[1] = 0;
 
 MPI_Cart_sub(GridComm, Subdims, &ColComm);
}

void BlockAComm(int i, double *BlockA, double* Block, int BSize)
{
 int tmp = (GridC[0] + i) % Grid;
 
 if (GridC[1] == tmp)
 {
  for (int j = 0; j < BSize*BSize; j++) BlockA[j] = Block[j];
 }
 
 MPI_Bcast(BlockA, BSize*BSize, MPI_DOUBLE, tmp, RowComm);
}

void BlockBComm(double *BlockB, int BSize)
{
 MPI_Status Status;
 int Next = GridC[0] + 1;
 if (GridC[0] == Grid - 1) Next = 0;
 int Prev = GridC[0] - 1;
 if (GridC[0] == 0) Prev = Grid - 1;

 MPI_Sendrecv_replace(BlockB, BSize*BSize, MPI_DOUBLE, Next, 0, Prev, 0, ColComm, &Status);
}

// последовательное перемножение матриц

void Mult(double* MatrixA, double* MatrixB, double* MatrixR, int Size)
{
  int i,j,k;   
  for (i = 0; i < Size; i++) 
    for (j = 0; j < Size; j++)
      for (k = 0; k < Size; k++)
        MatrixR[i*Size+j] += MatrixA[i*Size+k]*MatrixB[k*Size+j];
}

// метод Фокса

void Method(double* BlockA, double* Block, double* BlockB, double* BlockR, int BSize)
{
 for (int i = 0; i < Grid; i++)
 {
  BlockAComm(i, BlockA, Block, BSize);
  Mult(BlockA, BlockB, BlockR, BSize);
  BlockBComm(BlockB, BSize);
 }
}

// соединение данных

void CollectionResult(double *Result, double* BlockR, int Size, int BSize)
{
  double *Row = new double [Size*BSize];

  for (int i = 0; i < BSize; i++) MPI_Gather( &BlockR[i*BSize], BSize, MPI_DOUBLE, &Row[i*Size], BSize, MPI_DOUBLE, 0, RowComm);
 
  if (GridC[1] == 0) MPI_Gather(Row, BSize*Size, MPI_DOUBLE, Result, BSize*Size, MPI_DOUBLE, 0, ColComm);
  delete []Row;
}

int main(int argc, char* argv[])
{
 double *MatrixA = NULL, *BlockA = NULL; // матрица №1
 double *MatrixB = NULL, *BlockB = NULL; // матрица №2
 double *Result = NULL, *BlockR = NULL; // результат
 double *Block = NULL;
 int Size,BSize; // размер пространства

 double t1,t2;

 MPI_Init(&argc, &argv);
 
 MPI_Comm_size(MPI_COMM_WORLD, &Count);
 MPI_Comm_rank(MPI_COMM_WORLD, &Rank);

 Grid = sqrt((double)Count);

 if (Count != Grid*Grid) { if (Rank == 0) printf("Xren\n"); }
 else
 {
  if (Rank == 0) printf("Matrix multiplication: \n");

  CreateGird();

  if (Rank == 0)
  {
   do
   {
    printf("Enter size: ");
    scanf("%d", &Size);
 
    if (Size%Grid != 0) printf ("...\n");
   }
   while (Size%Grid != 0);
  }

  MPI_Bcast(&Size, 1, MPI_INT, 0, MPI_COMM_WORLD);

  BSize = Size / Grid;

  BlockA = new double [BSize*BSize];
  BlockB = new double [BSize*BSize];
  BlockR = new double [BSize*BSize];
  Block = new double [BSize*BSize];

  for (int i =0; i < BSize*BSize; i++) BlockR[i] = 0;

  if (Rank == 0)
  {
   MatrixA = new double [Size*Size];
   MatrixB = new double [Size*Size];
   Result = new double [Size*Size];
   InitializationMatrix(MatrixA, MatrixB, Size);
  }

  t1 = MPI_Wtime();

  Distribution(MatrixA,MatrixB,Block,BlockB,Size,BSize);

  Method(BlockA, Block, BlockB, BlockR, BSize);

  CollectionResult(Result, BlockR, Size, BSize);

  t2 = MPI_Wtime();

  if (Rank == 0) printf("Parallel method: %f\n",t2-t1);

  t1 = MPI_Wtime();

  Mult(MatrixA, MatrixB, Result, Size);

  t2 = MPI_Wtime();

  if (Rank == 0) printf("Sequential method: %f\n",t2-t1);

  delete[] MatrixA;
  delete[] MatrixB;
  delete[] Result; 
  delete[] BlockA;
  delete[] BlockB;
  delete[] BlockR;
 }

 MPI_Finalize();

 _getch();

 return 0;
}