// mpic++ -std=c++11 -Wall -Wextra -lboost_mpi -lboost_serialization -o laplacien_mpi.out laplacien_mpi.cpp 
// mpirun -n 50 --hostfile mpi_file ./laplacien_mpi.out

#include "image.hpp"

#include <boost/mpi.hpp>
namespace mpi = boost::mpi;

#include <algorithm>

int main(int argc, char ** argv)
{
    mpi::environment env(argc, argv);
    mpi::communicator world;

    int rank = world.rank();
    int size = world.size();
    int Nj = size;

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

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

        // some data for MPI transfers
        image_t data(sizeN);
        std::vector<image_t> allData(Nj, data);
        std::vector<image_t> allResults(Nj, data);

        // decompose image 
        for (int j=0; j<Nj; j++)
            std::copy_n(data1.begin()+(j*sizeN), sizeN,
                    std::inserter(allData[j], allData[j].begin()));

        // send subimage sizes to slave nodes
        mpi::broadcast(world, width, 0);
        mpi::broadcast(world, heightN, 0);

        // send data to slave nodes
        image_t nodeData;
        mpi::scatter(world, allData, nodeData, 0);

        // compute master data
        image_t nodeResult = computeLaplacian(nodeData, width, heightN, 10.0);

        // receive results from slave nodes
        mpi::gather(world, nodeResult, allResults, 0);

        // recompose image
        image_t data2(width*height2);
        for (int j=0; j<Nj; j++)
            std::copy_n(allResults[j].begin(), sizeN,
                    std::inserter(data2, data2.begin()+(j*sizeN)));

        // write output image
        writePgm("output.pgm", width, height, data2);
    }
    else  // slave nodes
    {
        // receive subimage sizes
        int width;
        mpi::broadcast(world, width, 0);
        int heightN;
        mpi::broadcast(world, heightN, 0);

        // receive data from master node
        image_t nodeData(width,heightN);
        mpi::scatter(world, nodeData, 0);

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

        // send results to master node
        mpi::gather(world, nodeResult, 0);
    }

    return 0;
}