// mpic++ -std=c++11 -Wall -Wextra -o laplacien_mpi_1.out laplacien_mpi_1.cpp 
// mpirun -n 4 ./laplacien_mpi_1.out

#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);
    double t0 = MPI_Wtime();

    // read image (master node)
    image_t data0;
    int width, height;
    if (worldRank == 0)  
    {
        std::string readError = readPgm("backloop.pgm", width, height, data0);
        if (readError != "")
        {
            std::cout << readError << std::endl;
            exit(-1);
        }
    }

    // broadcast image sizes
    MPI_Bcast(&width, 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(&height, 1, MPI_INT, 0, MPI_COMM_WORLD);

    int heightN = (height+worldSize-1) / worldSize;
    int sizeN = heightN * width;
    int height2 = heightN * worldSize;
    int size2 = height2 * width;

    // ensure height of data is a multiple of worldSize (master node)
    image_t data1(size2);
    if (worldRank == 0)
    {
        std::copy_n(data0.begin(), width*height, data1.begin());
    }

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

    // compute on each node
    image_t nodeResult = computeLaplacian(nodeData, width, heightN, 10.0);

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

    // write output image (master node)
    if (worldRank == 0)
    {
        writePgm("output_1.pgm", width, height, data2);
        double t1 = MPI_Wtime();
        std::cout << "walltime = " << t1 - t0 << std::endl;
    }

    MPI_Finalize();
    return 0;
}