#include <octave/oct.h>
#include "mpi.h"

#define MASTER 0/* taskid of first task */
#define FROM_MASTER 1/* setting a message type */
#define FROM_WORKER 2/* setting a message type */

MPI_Status status;

DEFUN_DLD (mpi_mm,args , ,"Parallel multiplication of two matrices using mpi")
{
  int argc;
  char **argv;

  Matrix a = args(0).matrix_value();
  Matrix b = args(1).matrix_value();
  Matrix c;

  int  NRA = a.rows();
  int  NCA = a.columns();
  int  NCB = b.columns();

  int numtasks,/* number of tasks in partition */
    taskid,/* a task identifier */
    numworkers,/* number of worker tasks */
    source,/* task id of message source */
    dest,/* task id of message destination */
    //    nbytes,/* number of bytes in message */
    mtype,/* message type */
    intsize,/* size of an integer in bytes */
    dbsize,/* size of a double float in bytes */
    rows,                      /* rows of matrix A sent to each worker */
    averow, extra, offset,      /* used to determine rows sent to each worker */
    i, j, k,/* misc */
    count;

  c(NRA, NCB);
  argc = 1;
  argv[0] = "4";

  intsize = sizeof(int);
  dbsize = sizeof(double);

  MPI_Init(&argc, &argv);
MPI_Comm_rank(MPI_COMM_WORLD, &taskid);
  MPI_Comm_size(MPI_COMM_WORLD, &numtasks);
  numworkers = numtasks-1;

  /**************************** master task ************************************/
  if (taskid == MASTER) {

    /* send matrix data to the worker tasks */
    averow = NRA/numworkers;
    extra = NRA%numworkers;
    offset = 0;
    mtype = FROM_MASTER;
    for (dest=1; dest<=numworkers; dest++) {
      rows = (dest <= extra) ? averow+1 : averow;

      MPI_Send(&offset, 1, MPI_INT, dest, mtype, MPI_COMM_WORLD);
      MPI_Send(&rows, 1, MPI_INT, dest, mtype, MPI_COMM_WORLD);
      count = rows*NCA;
      MPI_Send(&a(offset,0), count, MPI_DOUBLE, dest, mtype, MPI_COMM_WORLD);
      count = NCA*NCB;
      MPI_Send(&b, count, MPI_DOUBLE, dest, mtype, MPI_COMM_WORLD);

      offset = offset + rows;
    }

    /* wait for results from all worker tasks */
    mtype = FROM_WORKER;
    for (i=1; i<=numworkers; i++){
      source = i;
      MPI_Recv(&offset, 1, MPI_INT, source, mtype, MPI_COMM_WORLD, &status);
      MPI_Recv(&rows, 1, MPI_INT, source, mtype, MPI_COMM_WORLD, &status);
      count = rows*NCB;
      MPI_Recv(&c(offset,0), count, MPI_DOUBLE, source, mtype, MPI_COMM_WORLD,
               &status);
    }
    octave_value_list retval;


  retval(0) = c;

  return retval;


    //return c;//octave_value (c);
 }  /* end of master section */

  /**************************** worker task ************************************/
  if (taskid > MASTER) {
    mtype = FROM_MASTER;
    source = MASTER;

    MPI_Recv(&offset, 1, MPI_INT, source, mtype, MPI_COMM_WORLD, &status);

    MPI_Recv(&rows, 1, MPI_INT, source, mtype, MPI_COMM_WORLD, &status);

    count = rows*NCA;
    MPI_Recv(&a, count, MPI_DOUBLE, source, mtype, MPI_COMM_WORLD, &status);
    count = NCA*NCB;
    MPI_Recv(&b, count, MPI_DOUBLE, source, mtype, MPI_COMM_WORLD, &status);

    for (k=0; k<NCB; k++)
      for (i=0; i<rows; i++) {
        c(i,k) = 0.0;
        for (j=0; j<NCA; j++)
          c(i,k) = c(i,k) + a(i,j) * b(j,k);
      }

    mtype = FROM_WORKER;

    MPI_Send(&offset, 1, MPI_INT, MASTER, mtype, MPI_COMM_WORLD);
    MPI_Send(&rows, 1, MPI_INT, MASTER, mtype, MPI_COMM_WORLD);
    MPI_Send(&c, rows*NCB, MPI_DOUBLE, MASTER, mtype, MPI_COMM_WORLD);


  }  /* end of worker */
  MPI_Finalize();
} /* of main */