#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <mpi.h>
#define comm MPI_COMM_WORLD

void printMatrix(int nrows, int ncols, double *m) {
  for (int i = 0; i < nrows; i++) {
    for (int j = 0; j < ncols; j++) 
      printf("%f ", m[i*ncols + j]);
    printf("\n");
  }
  printf("\n");
}

void manager(int l, int m, int n, double *a, double *b, double *c) {
  double buf[m];
  int count, nprocs;
  MPI_Status status;
  
  MPI_Comm_size(comm, &nprocs);
  MPI_Bcast(b, l*m, MPI_DOUBLE, 0, comm);
  for (count = 0; count < nprocs-1 && count < n; count++)
    MPI_Send(a+count*l, l, MPI_DOUBLE, count+1, count+1, comm);
  for (int i = 0; i < n; i++) {
    MPI_Recv(buf, m, MPI_DOUBLE, MPI_ANY_SOURCE, MPI_ANY_TAG, comm, &status);
    memcpy(c+m*(status.MPI_TAG-1), buf, m*sizeof(double));
    if (count < n) {
      MPI_Send(a+count*l, l, MPI_DOUBLE, status.MPI_SOURCE, 1+count, comm);
      count++;
    }
  }
  for (int i = 1; i < nprocs; i++)
    MPI_Send(NULL, 0, MPI_INT, i, 0, comm);
}

void worker(int l, int m, int n) {
  double b[l*m], in[l], out[m];
  MPI_Status status;
    
  MPI_Bcast(b, l*m, MPI_DOUBLE, 0, comm);
  while (1) {
    MPI_Recv(in, l, MPI_DOUBLE, 0, MPI_ANY_TAG, comm, &status);
    if (status.MPI_TAG == 0) break;
    for (int j = 0; j < m; j++) {
      out[j] = 0.0;
      for (int k = 0; k < l; k++)
	out[j] += in[k]*b[k*m+j];
    }
    MPI_Send(out, m, MPI_DOUBLE, 0, status.MPI_TAG, comm);
  }
}

void mat_mul(int l, int m, int n, double *a, double *b, double *c) {
  int rank;

  MPI_Comm_rank(comm, &rank);
  if (rank == 0)
    manager(l, m, n, a, b, c);
  else
    worker(l, m, n);
}
