#include <stdio.h>
#include <stdlib.h>
//#include <mpi.h>
//#include <opencv2/opencv.hpp>

int main(int argc, char **argv)
{
    // Initialize MPI
    MPI_Init(&argc, &argv);

    // Get rank and size
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    // The input image
    cv::Mat image;

    // The total size of the image matrix (rows * columns * channels)
    size_t imageTotalSize = 0;

    // Partial size (how many bytes will be sent to each process)
    size_t imagePartialSize = 0;

    // How many channels are there in the image?
    int channels = 0;

    // Partial buffer, to contain the image
    uchar *partialBuffer = nullptr;

    // Also create the output image, where we will save the results
    cv::Mat outImage;

    // Read the image and its properties in the ROOT process
    if (rank == 0)
    {
        // Read the image
        image = cv::imread("image.jpg", cv::IMREAD_UNCHANGED);

        // Check if it's empty
        if (image.empty())
        {
            fprintf(stderr, "Image is empty, terminating!\n");
            MPI_Finalize();
            return -1;
        }

        // Get the number of channels in the image
        channels = image.channels();

        // Get the total size of the image matrix (rows * columns * channels)
        imageTotalSize = image.total() * image.elemSize();

        // Check if we can evenly divide the image bytes by the number of processes
        if (imageTotalSize % size)
        {
            fprintf(stderr, "Cannot evenly divide the image between the processes. Choose a different number of processes!\n");
            MPI_Finalize();
            return -2;
        }

        // Get partial size (how many bytes are sent to each process)
        imagePartialSize = imageTotalSize / size;

        printf("The image will be divided into blocks of %zu bytes each\n", imagePartialSize);
    }

    // Broadcast the "partial size" and the number of channels to all other processes
    MPI_Bcast(&imagePartialSize, 1, MPI_UNSIGNED_LONG_LONG, 0, MPI_COMM_WORLD);
    MPI_Bcast(&channels, 1, MPI_INT, 0, MPI_COMM_WORLD);

    // Synchronize the processes to make sure that the sizes are initialized
    MPI_Barrier(MPI_COMM_WORLD);

    // Allocate the partial buffer
    partialBuffer = (uchar *)malloc(imagePartialSize * sizeof(uchar));

    // Synchronize the processes to make sure each process has allocated the buffer
    MPI_Barrier(MPI_COMM_WORLD);

    // Scatter the image between the processes
    MPI_Scatter(image.data, imagePartialSize, MPI_UNSIGNED_CHAR, partialBuffer, imagePartialSize, MPI_UNSIGNED_CHAR, 0, MPI_COMM_WORLD);

    // Synchronize the image processing
    MPI_Barrier(MPI_COMM_WORLD);

    // Process the image
    for (size_t i = 0; i < imagePartialSize; i += channels)
    {
        // Get the pixel
        uchar *B = &partialBuffer[i];
        uchar *G = &partialBuffer[i + 1];
        uchar *R = &partialBuffer[i + 2];

        // Swap the blue and the red
        uchar temp = *B;
        *B = *R;
        *R = temp;
    }

    // Synchronize the image processing
    MPI_Barrier(MPI_COMM_WORLD);

    // Initialize the output image (only need to do it in the ROOT process)
    if (rank == 0)
    {
        outImage = cv::Mat(image.size(), image.type());
    }

    // Gather the processed image data to the ROOT process
    MPI_Gather(partialBuffer, imagePartialSize, MPI_UNSIGNED_CHAR, outImage.data, imagePartialSize, MPI_UNSIGNED_CHAR, 0, MPI_COMM_WORLD);

    // Save and display image (only in the ROOT process)
    if (rank == 0)
    {
        // Save the image
        cv::imwrite("new_image.jpg", outImage);

        // Show it on the screen
        cv::imshow("image", outImage);
        cv::waitKey(0);
        cv::destroyAllWindows();
    }

    // Clean up
    free(partialBuffer);
    MPI_Finalize();

    return 0;
}