/* julia.c */

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

/* size of the image width and height in pixels */
#define SIZE 5000

struct Pixel {
    unsigned char r;
    unsigned char g;
    unsigned char b;
};

/* an image is a 2D array of color data */
struct Image {
    int width, height;
    struct Pixel** pixels;
};

void save_pixels(FILE* file, struct Pixel* row, int width) {
    for (int i = 0; i < width; i++)  {
        fprintf(file, "%u ", row[i].r);
        fprintf(file, "%u ", row[i].g);
        fprintf(file, "%u ", row[i].b);
    }
    fprintf(file, "\n");

}

/* function to save the image to a PPM file */
void save_image(struct Image* image, const char* fname, int rank, int num_processors) {
    /* processor 0 does most of the work */
    if (rank == 0) {
        FILE* file = fopen(fname, "w");
        if (!file) {
            fprintf(stderr, "Error, could not open file '%s' for writing.\n", fname);
            return;
        }

        /* write the magic number, size, and max color */
        fprintf(file, "P3\n%d %d\n255\n", SIZE, SIZE);
        
        /* write our own data */
        for (int i = 0; i < image->height; i++) {
            save_pixels(file, image->pixels[i], image->width);
        }
        
        /* for each other process, receive it's data and write it */
        struct Pixel* row = malloc(sizeof(struct Pixel) * image->width);
        for (int i = 1; i < num_processors; i++) {
            for (int j = 0; j < image->height; j++) {
                MPI_Recv(row, image->width * sizeof(struct Pixel), MPI_BYTE, i, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                save_pixels(file, row, image->width);
            }
        }

        /* close the file */
        fclose(file);
    } else {
        /* we're another process, just send the data to process 0 */
        for (int i = 0; i < image->height; i++) {
            MPI_Send(image->pixels[i], image->width * sizeof(struct Pixel), MPI_BYTE, 0, 0, MPI_COMM_WORLD);
        }
    }
}

/* a struct for storing complex numbers */
struct Complex {
    float r;
    float i;
};

/* get the magnitude of a complex number */
float magnitude(struct Complex c) {
    return c.r * c.r + c.i * c.i;
}

/* multiply two complex numbers */
struct Complex mult(struct Complex a, struct Complex b) {
    struct Complex result;
    result.r = a.r * b.r - a.i * b.i;
    result.i = a.i * b.r + a.r * b.i;
    return result;
}

/* add two complex numbers */
struct Complex add(struct Complex a, struct Complex b) {
    struct Complex result;
    result.r = a.r + b.r;
    result.i = a.i + b.i;
    return result;
}

/* the julia function */
int julia(int x, int y) {
    const float scale = 1.5;
    float jy = scale * (float)(SIZE / 2 - x) / (SIZE / 2);
    float jx = scale * (float)(SIZE / 2 - y) / (SIZE / 2);

    struct Complex c;
    c.r = -0.8;
    c.i = 0.156;

    struct Complex a;
    a.r = jx;
    a.i = jy;

    int i;
    for (i = 0; i < 200; i++) {
        a = add(mult(a, a), c);
        if (magnitude(a) > 1000) {
            return 0;
        }
    }

    return 1;
}

/* the main kernel function runs julia on each pixel */
void generateImage(struct Image* image, int rank, int num_processors) {
    int start = (SIZE / num_processors) * rank;
    int end = start + (SIZE / num_processors);

    for (int row = start; row < end; row++) {
        for (int col = 0; col < image->width; col++) {
            /* calculate the value at this position */
            int julia_value = julia(row, col);
            
            /* use this to weight the red amount */
            image->pixels[row - start][col].r = 255 * julia_value;
            image->pixels[row - start][col].g = 0;
            image->pixels[row - start][col].b = 0;
        }
    }
}

/* the main function begins running on the CPU */
int main(int argc, char** argv) {
    MPI_Init(&argc, &argv);

    /* get our rank and size */
    int rank, num_processors;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &num_processors);
    
    /* create the image for just our portion */
    struct Image image;
    image.width = SIZE;
    image.height = SIZE / num_processors;
    image.pixels =  malloc(sizeof(struct Pixel*) * image.height);
    for (int i = 0; i < image.height; i++) {
        image.pixels[i] = malloc(sizeof(struct Pixel) * image.width);
    }

    /* generate the image in parallel */
    generateImage(&image, rank, num_processors);
   
    /* save the image */
    save_image(&image, "julia.ppm", rank, num_processors); 

    MPI_Finalize();

    return 0;
}