// mpic++ -std=c++11 -Wall -Wextra -o laplacien_mpi_2.out laplacien_mpi_2.cpp 
// mpirun -n 4 ./laplacien_mpi_2.out canyon.pgm canyon_mpi_2.pgm
// separate master code and slave code 

#include "image.hpp"
#include <algorithm>
#include <mpi.h>

int main(int argc, char ** argv)
{
    // init MPI
    MPI_Init(&argc, &argv);
    int worldSize;
    MPI_Comm_size(MPI_COMM_WORLD, &worldSize);
    int worldRank;
    MPI_Comm_rank(MPI_COMM_WORLD, &worldRank);

    if (argc != 5)
    {
        std::cout << "usage: " << argv[0] 
            << " <input> <output> <scaling> <nb fakes>\n";
        exit(-1);
    }
    const char * INPUT = argv[1];
    const char * OUTPUT = argv[2];
    const float SCALING = atof(argv[3]);
    const int NB_FAKES = atoi(argv[4]);

    if (worldRank == 0)  // master node
    {
        double t0 = MPI_Wtime();

        // read image
        image_t data0;
        int width, height;
        std::string readError = readPgm(INPUT, width, height, data0);
        if (readError != "")
        {
            std::cout << readError << std::endl;
            exit(-1);
        }
        int heightN = (height+worldSize-1) / worldSize;
        int sizeN = heightN * width;
        int height2 = heightN * worldSize;
        int size2 = height2 * width;

        // copy data0 in data1 
        // ensure the height of data1 is a multiple of worldSiz
        image_t data1(size2);
        std::copy_n(data0.begin(), width*height, data1.begin());

        // send subimage sizes to slave nodes
        MPI_Bcast(&width, 1, MPI_INT, 0, MPI_COMM_WORLD);
        MPI_Bcast(&heightN, 1, MPI_INT, 0, MPI_COMM_WORLD);

        // send data to slave nodes
        image_t nodeData(sizeN);
        MPI_Scatter(data1.data(), sizeN, MPI_UNSIGNED_CHAR, 
                nodeData.data(), sizeN, MPI_UNSIGNED_CHAR, 
                0, MPI_COMM_WORLD);

        // compute master data
        image_t nodeResult = computeLaplacian(nodeData, width, heightN, SCALING);
        for (int k=0; k<NB_FAKES; ++k)
            nodeResult = computeLaplacian(nodeData, width, heightN, SCALING);

        // receive results from slave nodes
        image_t data2(size2, 0);
        MPI_Gather(nodeResult.data(), sizeN, MPI_UNSIGNED_CHAR, 
                data2.data(), sizeN, MPI_UNSIGNED_CHAR, 
                0, MPI_COMM_WORLD);

        // write output image
        writePgm(OUTPUT, width, height, data2);
        double t1 = MPI_Wtime();
        std::cout << t1 - t0;
    }
    else  // slave nodes
    {
        // receive subimage sizes
        int width;
        MPI_Bcast(&width, 1, MPI_INT, 0, MPI_COMM_WORLD);
        int heightN;
        MPI_Bcast(&heightN, 1, MPI_INT, 0, MPI_COMM_WORLD);
        int sizeN = heightN * width;

        // receive data from master node
        image_t nodeData(sizeN);
        MPI_Scatter(nullptr, sizeN, MPI_UNSIGNED_CHAR, 
                nodeData.data(), sizeN, MPI_UNSIGNED_CHAR, 
                0, MPI_COMM_WORLD);

        // compute node data
        image_t nodeResult = computeLaplacian(nodeData, width, heightN, SCALING);
        for (int k=0; k<NB_FAKES; ++k)
            nodeResult = computeLaplacian(nodeData, width, heightN, SCALING);

        // send results to master node
        MPI_Gather(nodeResult.data(), sizeN, MPI_UNSIGNED_CHAR, 
                nullptr, sizeN, MPI_UNSIGNED_CHAR, 
                0, MPI_COMM_WORLD);
    }

    MPI_Finalize();
    return 0;
}