/* julia.cu */

#include <stdio.h>

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

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

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

    /* write the magic number, size, and max color */
    fprintf(f, "P3\n%d %d\n255\n", image->width, image->height);

    /* write the image data */
    int i, j;
    for (i = 0; i < image->height; i++) {
        for (j = 0; j < image->width; j++)  {
            fprintf(f, "%u ", image->data[i*image->width*3 + j*3 + 0]);
            fprintf(f, "%u ", image->data[i*image->width*3 + j*3 + 1]);
            fprintf(f, "%u ", image->data[i*image->width*3 + j*3 + 2]);
        }
        fprintf(f, "\n");
    }

    /* close the file */
    fclose(f);
}

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

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

/* multiply two complex numbers */
__device__ 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 */
__device__ 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 */
__device__ int julia(int x, int y) {
    const float scale = 1.5;
    float jx = scale * (float)(SIZE / 2 - x) / (SIZE / 2);
    float jy = 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 */
__global__ void kernel(unsigned char* data) {
    /* map from blockIdx to pixel position */
    int x = blockIdx.x;
    int y = blockIdx.y;

    int offset = x + y * gridDim.x;

    /* calculate the value at this position */
    int julia_value = julia(x, y);
    data[offset * 3 + 0] = 255 * julia_value;
    data[offset * 3 + 1] = 0;
    data[offset * 3 + 2] = 0;
}

/* the main function begins running on the CPU */
int main(int argc, char** argv) {
    /* create the image */
    struct Image image;
    image.width = SIZE;
    image.height = SIZE;
    image.data = (unsigned char*) malloc(SIZE * SIZE * 3 * sizeof(unsigned char));

    /* allocate space for the image data */
    unsigned char* gpu_image;
    cudaMalloc((void**) &gpu_image, SIZE * SIZE * 3 * sizeof(unsigned char));

    /* run the kernel */
    dim3 grid(SIZE, SIZE);
    kernel<<<grid, 1>>>(gpu_image);

    /* copy the data back into the image */
    cudaMemcpy(image.data, gpu_image, SIZE * SIZE * 3 * sizeof(unsigned char), cudaMemcpyDeviceToHost);

    /* free the GPU memory */
    cudaFree(gpu_image);

    /* save the image */
    save_image(&image, "julia.ppm"); 

    return 0;
}