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

#define MASTER_RANK 0
#define BLUR_RADIUS 1

// Function to blur a pixel
void blurPixel(unsigned char *image, int width, int x, int y) {
    int sum = 0;
    int count = 0;
    for (int i = x - BLUR_RADIUS; i <= x + BLUR_RADIUS; i++) {
        for (int j = y - BLUR_RADIUS; j <= y + BLUR_RADIUS; j++) {
            if (i >= 0 && i < width && j >= 0 && j < width) {
                sum += image[i * width + j];
                count++;
            }
        }
    }
    image[x * width + y] = (unsigned char)(sum / count);
}

// Function to blur a block of the image
void blurBlock(unsigned char *block, int width, int block_size) {
    for (int i = 0; i < block_size; i++) {
        for (int j = 0; j < block_size; j++) {
            blurPixel(block, width, i, j);
        }
    }
}

int main(int argc, char *argv[]) {
    int rank, size;
    int image_size = 512;
    int block_size = 64;
    int num_blocks = image_size / block_size;

    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    // Check if number of processes is valid
    if (size != num_blocks * num_blocks) {
        if (rank == MASTER_RANK) {
            printf("Number of processes must equal the number of image blocks\n");
        }
        MPI_Finalize();
        return EXIT_FAILURE;
    }

    // Generate image data on the master process
    unsigned char *image = NULL;
    if (rank == MASTER_RANK) {
        image = (unsigned char *)malloc(image_size * image_size * sizeof(unsigned char));
        // Initialize image with random data
        for (int i = 0; i < image_size * image_size; i++) {
            image[i] = rand() % 256;
        }
    }

    // Scatter image blocks to processes
    unsigned char *block = (unsigned char *)malloc(block_size * block_size * sizeof(unsigned char));
    MPI_Scatter(image, block_size * block_size, MPI_UNSIGNED_CHAR, block, block_size * block_size, MPI_UNSIGNED_CHAR, MASTER_RANK, MPI_COMM_WORLD);

    // Process the block
    blurBlock(block, block_size, block_size);

    // Gather processed blocks back to master process
    MPI_Gather(block, block_size * block_size, MPI_UNSIGNED_CHAR, image, block_size * block_size, MPI_UNSIGNED_CHAR, MASTER_RANK, MPI_COMM_WORLD);

    // Print the processed image
    if (rank == MASTER_RANK) {
        printf("Processed Image:\n");
        for (int i = 0; i < image_size; i++) {
            for (int j = 0; j < image_size; j++) {
                printf("%3d ", image[i * image_size + j]);
            }
            printf("\n");
        }
    }

    // Cleanup
    free(image);
    free(block);
    MPI_Finalize();

    return EXIT_SUCCESS;
}