/* 
 *  Generic example of master/worker program with dynamic assignment of
 *  tasks to workers.  (The point of this example is to demonstrate use
 *  of tags.)  "Tasks" just consist of sleeping for a specified
 *  time; times to wait are generatedly randomly from 0 to a specified
 *  maximum.
 *
 *  Command-line arguments:  total tasks, maximum wait time
 */
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <mpi.h>

#define WORK_TAG 0              /* "do this work" messages */
#define STOP_TAG 1              /* "stop worker" messages */
#define WORK_DONE_TAG 2         /* "work done" messages */
#define RESULT_TAG 3            /* "result" messages */

/* a short function to print a message, stop all processes, and exit */
void error_exit(char msg[]) {
    fprintf(stderr, "%s", msg);
    MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
}

/* function to generate a random integer in [start, end] */
int random_in_range(int start, int end) {
    /* use method described in "rand" man page */
    return start + ((double) (end-start+1) * rand()/(RAND_MAX + 1.0));
}


/* ---- Program for master process ----------------------------------- */

void master_process (int nworkers, int argc, char *argv[]) {

    int total_tasks = 0;
    int max_wait = 0;
    int tasks_done = 0;
    int work_amt;
    int i;
    int workerID;
    int worker_result[2];
    MPI_Status status;
    int min_count, max_count, min_time, max_time; 

    if (argc < 3) 
        error_exit("arguments are total_tasks, max_wait\n");
    total_tasks = atoi(argv[1]);
    max_wait = atoi(argv[2]);
    if ((total_tasks <= 0) || (max_wait <= 0)) 
        error_exit("arguments are total_tasks, max_wait\n");

    printf("%d workers, %d total tasks, %d maximum wait\n", 
            nworkers, total_tasks, max_wait);

    /* Give each worker one piece of work initially */
    for (i = 0; i < nworkers && i < total_tasks; ++i) {
        work_amt = random_in_range(0, max_wait);
        MPI_Send(&work_amt, 1, MPI_INT, i+1, WORK_TAG, MPI_COMM_WORLD);
    }

    /* While not all work done .... */
    while (tasks_done < total_tasks) {

        /* Receive from worker */
        MPI_Recv(&work_amt, 1, MPI_INT, MPI_ANY_SOURCE, WORK_DONE_TAG, 
                MPI_COMM_WORLD, &status);
        ++tasks_done;
        workerID = status.MPI_SOURCE;

        /* If more work to do, assign to worker that just finished */
        if (tasks_done < total_tasks) {
            work_amt = random_in_range(0, max_wait);
            MPI_Send(&work_amt, 1, MPI_INT, workerID, WORK_TAG, MPI_COMM_WORLD);
        }
      }

    /* Shut down workers */
    for (i = 0; i < nworkers; ++i) {
        MPI_Send(&work_amt, 0, MPI_INT, i+1, STOP_TAG, MPI_COMM_WORLD);
    }

    /* Get their results and print */
    for (i = 0; i < nworkers; ++i) {
        MPI_Recv(worker_result, 2, MPI_INT, i+1, RESULT_TAG, 
                MPI_COMM_WORLD, &status);
        printf("worker %d did %d tasks, total time = %d\n", 
                i+1, worker_result[0], worker_result[1]);
        if (i == 0) {
            min_count = max_count = worker_result[0];
            min_time = max_time = worker_result[1];
        }
        else {
            if (worker_result[0] < min_count) min_count = worker_result[0];
            if (worker_result[0] > max_count) max_count = worker_result[0];
            if (worker_result[1] < min_time) min_time = worker_result[1];
            if (worker_result[1] > max_time) max_time = worker_result[1];
        }
    }
    printf("task count ranged from %d to %d\n", min_count, max_count);
    printf("time ranged from %d to %d\n", min_time, max_time);
}

/* ---- Program for worker process ---------------------------------- */

void slave_process (int myID) {

    MPI_Status status;
    int work_amt;
    int task_count = 0;
    int total_work_time = 0;
    int buffer[2];

    /* While master is still sending "do this work" messages .... */
    while ((MPI_Recv(&work_amt, 1, MPI_INT, 0, MPI_ANY_TAG, MPI_COMM_WORLD,
                    &status) == MPI_SUCCESS) 
            && (status.MPI_TAG == WORK_TAG)) {

        /* Get a message and fake work */
        printf("worker %d received %d\n", myID, work_amt);
        ++task_count;
        total_work_time += work_amt;
        sleep(work_amt);
        MPI_Send(&work_amt, 1, MPI_INT, 0, WORK_DONE_TAG, MPI_COMM_WORLD);
    }
    printf("worker %d done\n", myID);
    buffer[0] = task_count;
    buffer[1] = total_work_time;
    MPI_Send(buffer, 2, MPI_INT, 0, RESULT_TAG, MPI_COMM_WORLD);
}

/* ---- Main program ------------------------------------------------ */

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

  int nprocs;
  int myID;

  /* initialize for MPI; get number of processes and my ID */

  if (MPI_Init(&argc, &argv) != MPI_SUCCESS) {
      fprintf(stderr, "MPI initialization error\n"); 
      exit(EXIT_FAILURE);
  }

  MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
  MPI_Comm_rank(MPI_COMM_WORLD, &myID);
  if (nprocs < 2) {
      fprintf(stderr, "Must have at least two processes\n");
      MPI_Finalize();
      exit(EXIT_FAILURE);
  }

  /* start one master process and nprocs-1 workers */

  if (myID == 0) 
      master_process(nprocs-1, argc, argv);
  else 
      slave_process(myID);

  /* clean up and return */

  MPI_Finalize();
  return EXIT_SUCCESS;
}