/*201220 Roberto Giorgi - University of Siena, Italy
##########################################################################
# MIT LICENSE:
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHERi
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
##########################################################################*/

#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>

#define HISTOGRAM_BIN_COUNT 4
char* Color;
int* Histogram;
#define N 32

#define MASTER 0
#define TAG_GENERAL 1
void histo_mpi(int *histogram, char *color, int size, int wid, int nw) {
    int rc, i, histogram_private[HISTOGRAM_BIN_COUNT];
    for(i=0; i<HISTOGRAM_BIN_COUNT; i++) histogram_private[i] = 0;

    int ssize = (size - 1) /nw + 1;
    char *slice = malloc(ssize * sizeof(char));
    if (slice == NULL) { 
        printf("Worker %d: ERROR while allocating 'slice'.\n", wid); fflush(stdout);
        exit(1);
    }

    if (wid == MASTER) { // Distribute the work
        int w, i;

        // Assign the first slice to the master
        for (i= 0; i<ssize; ++i) slice[i] = color[i];

        // Send the other slices to the slaves
        for(w=1; w<nw; ++w) {
            char *start = color + w*ssize;
            rc = MPI_Send(start, ssize, MPI_CHAR, w, TAG_GENERAL, MPI_COMM_WORLD);
            printf("Master(%d): sent message at %08lX to worker %d with tag %d and size %d.\n",
                MASTER, (unsigned long)start, w, TAG_GENERAL, ssize); fflush(stdout);
        }
    } else { // slave
        // Wait until a message is there to be received
        int dataWaitingFlag;
        printf("Worker %d is waiting...\n", wid); fflush(stdout);
        do {
            MPI_Iprobe(MASTER, TAG_GENERAL, MPI_COMM_WORLD, &dataWaitingFlag, MPI_STATUS_IGNORE);
        } while (!dataWaitingFlag);
        printf("Worker %d: message has arrived from master.\n", wid); fflush(stdout);
        
        // Get the message and put it in 'slice'
        MPI_Status Stat;
        rc = MPI_Recv(slice, ssize, MPI_CHAR, MASTER, TAG_GENERAL, MPI_COMM_WORLD, &Stat);

        // Verify correct size
        int count;
        rc = MPI_Get_count(&Stat, MPI_CHAR, &count);
        printf("Worker %d: received %d char(s) from master(%d) with tag %d \n",
            wid, count, Stat.MPI_SOURCE, Stat.MPI_TAG); fflush(stdout);
    }
    
    // Processing data
    printf("Worker %d: start processing...\n", wid); fflush(stdout);
    for (int i= 0; i<ssize; ++i) {
        histogram_private[slice[i]]++;
        printf("w%03d  i=%-3d  c=%-3d  d=%-d\n", wid, i, slice[i], histogram_private[slice[i]]); fflush(stdout);
    }
    free(slice);

    if (wid == MASTER) { // Process the partial results
        int done = 0, w;

        // Accumulate the result
        for (int i= 0; i<HISTOGRAM_BIN_COUNT; ++i) {
            for (int k=0; k < HISTOGRAM_BIN_COUNT; ++k) printf("%3d ", histogram_private[k]);
            histogram[i] += histogram_private[i];
            printf("w%03d  i=%-3d  d=%-d\n", wid, i, histogram[i]); fflush(stdout);
        }

        // Get partial histograms from slaves
        do {
            for (w=1; w<nw; ++w) { // round robin check
                int dataWaitingFlag;
                MPI_Iprobe(w, TAG_GENERAL, MPI_COMM_WORLD, &dataWaitingFlag, MPI_STATUS_IGNORE);
                if (dataWaitingFlag) {
                    printf("Master: data arrived from worker %d.\n", wid); fflush(stdout);

                    // Get the message and put it in 'histogram_private'
                    MPI_Status Stat;
                    rc = MPI_Recv(histogram_private, HISTOGRAM_BIN_COUNT, MPI_INT, w, TAG_GENERAL, MPI_COMM_WORLD, &Stat);
                    printf("Master has received data (rc=%d).\n", rc); fflush(stdout);
                    ++done;

                    // Verify correct size
                    int count;
                    rc = MPI_Get_count(&Stat, MPI_INT, &count);
                    printf("Master: received %d int(s) (%d) from worker %d with tag %d (done=%d)\n",
                        wid, count, Stat.MPI_SOURCE, Stat.MPI_TAG, done);

                    // Accumulate the result
                    for (int i= 0; i<HISTOGRAM_BIN_COUNT; ++i) {
                        for (int k=0; k < HISTOGRAM_BIN_COUNT; ++k) printf("%3d ", histogram_private[k]);
                        histogram[i] += histogram_private[i];
                        printf("w%03d  i=%-3d  d=%-d\n", wid, i, histogram[i]); fflush(stdout);
                    }
                }
            }
        } while (done < nw - 1);
    } else { // slave: send back the partial result
        rc = MPI_Send(histogram_private, HISTOGRAM_BIN_COUNT, MPI_INT, MASTER, TAG_GENERAL, MPI_COMM_WORLD);
        printf("Worker %d: sent message at %08lX to master(%d) with tag %d and size %d (rc=%d).\n",
            wid, (unsigned long)histogram_private , MASTER, TAG_GENERAL, ssize, rc); fflush(stdout);
    }
}

int main(int argc, char **argv) {
    // Initializing MPI environment
    int NumProcesses;
    int Rank;
    MPI_Init(&argc,&argv);

    // Gets number of tasks/processes that this program is running on
    MPI_Comm_size(MPI_COMM_WORLD, &NumProcesses);

    // Gets the rank (process/task number) that this program is running on
    MPI_Comm_rank(MPI_COMM_WORLD, &Rank);
    
    // Print initial data
    if (Rank == MASTER) {
        Color = (char*)malloc(N * sizeof(char));
        Histogram = (int*)malloc(HISTOGRAM_BIN_COUNT * sizeof(int));
        srand(2017);
        for (int i = 0; i < N; ++i) Color[i] = (char)(rand() % HISTOGRAM_BIN_COUNT);
        for (int i = 0; i < HISTOGRAM_BIN_COUNT; ++i) Histogram[i] = 0;
        for(int i=0; i<N; i++) printf("%3d ", Color[i]); printf("\n"); fflush(stdout);
    }

    ///////////////////////////////
    // Call the processing function
    histo_mpi(Histogram, Color, N, Rank, NumProcesses);
    ///////////////////////////////

    // Print final data
    if (Rank == MASTER) {
        printf("---\n"); fflush(stdout);
        for(int c=0; c<HISTOGRAM_BIN_COUNT; c++)
            printf("c=%3d  d=%d\n", c, Histogram[c]); fflush(stdout);
        free(Color); free(Histogram);
    }

    // Shut down MPI
    MPI_Finalize();
}