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

int main(int argc, char** argv) {

    int size, rank, a, b;
    MPI_Status status;
    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    a = rank;
    b = -1;
    int s;
    int* arr_B;
    int* arr_A;
    double start = 0, end = 0;
    if (rank == 0) {
        scanf_s("%d", &s);
        printf("%d\n", s);
        start = MPI_Wtime();
    }


        MPI_Bcast(&s, 1, MPI_INT, 0, MPI_COMM_WORLD);
        arr_A = new int[s * s];
        arr_B = new int[s * s];
        if (rank ==0 )
        {
            for (int i = 0; i < s; i++)
                for (int j = 0; j < s; j++) {
                    arr_A[i * s + j] = (rand() % 100 + 1);
                    arr_B[i * s + j] = (rand() % 100 + 1);

                }
        }
        MPI_Bcast(arr_B, s * s, MPI_INT, 0, MPI_COMM_WORLD);
        int part_size = s * s / size;
        std::cout << part_size;
        int* arr_row = new int[part_size];

        MPI_Scatter(arr_A, part_size, MPI_INT, arr_row, part_size, MPI_INT, 0, MPI_COMM_WORLD);
        int* arr_ans = new int[s * s / size];

        for (int i = 0; i < part_size / s; i++) {
            for (int j = 0; j < start; j++) {
                int sum = 0;
                for (int k = 0; k < s; k++)
                {
                    sum += arr_row[i * s + k] * arr_B[k * s + j];
                }
                arr_ans[i * s + j] = sum;
            }

        }

        MPI_Gather(arr_ans, part_size, MPI_INT, arr_A, part_size, MPI_INT, 0, MPI_COMM_WORLD);

        if (rank == 0) {
            end = MPI_Wtime();
            std::cout << "TIME:" << end - start;

        }
        delete[](arr_row);
        delete[](arr_ans);
        if (rank == 0) {
            delete[](arr_B);
            delete[](arr_A);
        }
        MPI_Barrier(MPI_COMM_WORLD);
        MPI_Finalize();


}