#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <math.h>
#include <sys/wait.h>
#include <sys/mman.h>

#include "images/images.h"
#include "complex.c"


#define MB_THREADS 11

typedef struct _mb_settings {
    unsigned int width;
    unsigned int iterations;
    double re_start;
    double re_end;
    double im_start;
    double im_end;
} MandelbrotSettings;

MandelbrotSettings mb_settings_default(int width, int iterations) {
    MandelbrotSettings set;
    set.width = width;
    set.iterations = iterations;
    set.re_start = -2;
    set.re_end = 1;
    set.im_start = -1.25;
    set.im_end = 1.25;
    return set;
}

MandelbrotSettings mb_settings_at_pt(int width, int iterations, double re, double im, double px_step) {
    MandelbrotSettings set;
    set.width = width;
    int height = width;
    set.iterations = iterations;
    set.re_start = re - (width * (0.5 *px_step));
    set.re_end = re + (width * (0.5 *px_step));
    set.im_start = im - (height * (0.5 *px_step));
    set.im_end = im + (height * (0.5 *px_step));
    return set;
}

void mb_print_settings(MandelbrotSettings set) {
    printf(
        "Settings: {\n\twidth; %d\n\titerations; %d\n\tre_start; %f\n\tre_end; %f\n\tim_start; %f\n\tim_end; %f\n}\n",
        set.width,
        set.iterations,
        set.re_start,
        set.re_end,
        set.im_start,
        set.im_end
    );
}

void draw_mandelbrot(MandelbrotSettings settings);
unsigned int test_mandelbrot(Complex c,  int iterations);
void* create_shared_memory(size_t size);
int image_new_shared(int width, int height, Image* img);
void image_destroy_shared(Image img);

int main(int argc, char* argv[]) {
    int width = 3000;
    int iterations = 60;
    if (argc > 1) {
        width = atoi(argv[1]);
    }
    if (argc > 2) {
        iterations = atoi(argv[2]);
    }

    double step = (1/ (double) 18) / width;

    MandelbrotSettings set = mb_settings_at_pt(
        width, iterations, -1, -(2.2/(double)9), step
    );

    mb_print_settings(set);

    draw_mandelbrot(set);
}


void draw_mandelbrot(MandelbrotSettings set) {
    double step = (set.re_end - set.re_start) / set.width;
    unsigned int height = (set.im_end - set.im_start) / step;

    printf("Dimensions: %ux%u (step: %f)\n", set.width, height, step);

    Image img;
    image_new_shared(set.width, height, &img);

    unsigned int thread_id = 0;

    for (unsigned int i = 0; i < MB_THREADS - 1; i++) {
        if (fork() == 0) {
            thread_id = i + 1;
            break;
        }
    }
    printf("Thread %i reporting for duty\n", thread_id);

    for (unsigned int x = 0; x < set.width; x++) {
        if (x % MB_THREADS != thread_id) continue;

        if (x % (set.width / 100) == 0) {
            printf("Rendering... %02i%%\r", (int) ((x / (float) set.width) * 100));
            fflush(stdout);
        }

        for (unsigned int y = 0; y < height; y++) {
            Complex c = complex_new(
                set.re_start + (x * step),
                set.im_start + (y * step)
            );
            unsigned int ret = test_mandelbrot(c, set.iterations);
            //int r = (ret * 160 / iterations);
            //int g = (ret * 9 / iterations);
            //int b = (ret * 165 / iterations);
            //image_set_px(img, x, y, r,g,b);
            int color = 255 * sqrt(ret / (float) set.iterations);
            image_set_px(img, x, y, color, color, color);
        }
    }

    printf("Thread %i is finished\n", thread_id);
    if (thread_id != 0) {
        exit(0);
    }

    int status;
    for (int i = 0; i < MB_THREADS - 1; i++) {
        while(wait(&status) > 0) {}
    }

    printf("Got %d/%d threads\n", MB_THREADS - 1, MB_THREADS - 1);

    image_save_bmp(img, "mandelbrot.bmp");
    image_destroy_shared(img);
}

unsigned int test_mandelbrot(Complex c,  int iterations) {
    Complex z = complex_new(0,0);

    for (int i = 0; i < iterations; i++) {
        z = complex_add(complex_mult(z,z), c); // z^2 + c
        if (complex_abs(z) > 2) return i;
    }
    return 0;
}



int image_new_shared(int width, int height, Image* img) {
    if (img == NULL) return 0;

    img->height = height;
    img->width = width;

    if ((unsigned long int) width * height * 3 > (unsigned long int) UINT_MAX) {
        printf("Image dimensions to large for a bitmap!\n");
        return 0;
    }

    // initialize bitmap...
    img->bitmap = create_shared_memory(height * sizeof(char*));
    for (int i = 0; i < height; i++) {
        img->bitmap[i] = create_shared_memory(3 * width * sizeof(char));
    }

    return 1;
}

void image_destroy_shared(Image img) {
    for (int i = 0; i < img.height; i++) {
        munmap(img.bitmap[i], 3 * img.width * sizeof(char));
    }  
    munmap(img.bitmap, img.height * sizeof(char*));
}

void* create_shared_memory(size_t size) {
  // Our memory buffer will be readable and writable:
  int protection = PROT_READ | PROT_WRITE;

  // The buffer will be shared (meaning other processes can access it), but
  // anonymous (meaning third-party processes cannot obtain an address for it),
  // so only this process and its children will be able to use it:
  int visibility = MAP_SHARED | MAP_ANONYMOUS;

  // The remaining parameters to `mmap()` are not important for this use case,
  // but the manpage for `mmap` explains their purpose.
  return mmap(NULL, size, protection, visibility, -1, 0);
}