#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <math.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>
#include <sys/time.h>
#include "mpi.h"

static
void init_array(int ni, int nj, int nk,
                float* alpha,
                float* beta,
                float C[ni][nj],
                float A[ni][nk],
                float B[nk][nj]) {
    int i, j;

    *alpha = 1.5;
    *beta = 1.2;
    for (i = 0; i < ni; i++)
        for (j = 0; j < nj; j++)
            C[i][j] = (float) ((i * j + 1) % ni) / ni;
    for (i = 0; i < ni; i++)
        for (j = 0; j < nk; j++)
            A[i][j] = (float) (i * (j + 1) % nk) / nk;
    for (i = 0; i < nk; i++)
        for (j = 0; j < nj; j++)
            B[i][j] = (float) (i * (j + 2) % nj) / nj;
}

static
void kernel_gemm(int ni, int nj, int nk,
                 float alpha,
                 float beta,
                 float C[ni][nj],
                 float A[ni][nk],
                 float B[nk][nj],
                 int i_1,
                 int i_2) {
    int i, j, k;

    for (i = i_1; i < i_2; i++) {
        for (j = 0; j < nj; j++) {
            C[i][j] *= beta;
            for (k = 0; k < nk; k++) {
                C[i][j] += alpha * A[i][k] * B[k][j];
            }
        }
    }
}

int main(int argc, char** argv) {
        int proc_num, myrank;
    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &proc_num);
    MPI_Comm_rank(MPI_COMM_WORLD, &myrank);

    int ni = atoi(argv[1]);
    int nj = atoi(argv[2]);
    int nk = atoi(argv[3]);
    if (myrank == 0) {
        printf("%d processes, i = %d, j = %d, k = %d, ", proc_num, ni, nj, nk);
    }
    int i_min = ni / proc_num * myrank;
    int i_max = ni / proc_num * (myrank + 1);
    if (myrank < ni % proc_num) {
        i_min += myrank;
        i_max += myrank + 1;
    } else {
        i_min += ni % proc_num;
        i_max += ni % proc_num;
    }

    float alpha;
    float beta;
    float (* C)[ni][nj];
    C = (float (*)[ni][nj]) malloc((ni) * (nj) * sizeof(float));
    float (* A)[ni][nk];
    A = (float (*)[ni][nk]) malloc((ni) * (nk) * sizeof(float));
    float (* B)[nk][nj];
    B = (float (*)[nk][nj]) malloc((nk) * (nj) * sizeof(float));

    init_array(ni, nj, nk, &alpha, &beta,
               *C,
               *A,
               *B);

    MPI_Barrier(MPI_COMM_WORLD);
    double time1 = MPI_Wtime();

    kernel_gemm(ni, nj, nk,
                alpha, beta,
                *C,
                *A,
                *B,
                i_min,
                i_max);

    MPI_Barrier(MPI_COMM_WORLD);
    double time2 = MPI_Wtime();
    if (myrank == 0) {
        printf(" %lf sec\n", time2 - time1);
    }


    free((void*) C);
    free((void*) A);
    free((void*) B);

    MPI_Finalize();
    return 0;
}