#include <assert.h>
#include <complex.h>
#include <getopt.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>
#include <png.h>
#include <pthread.h>

typedef uint32_t pixel_t;

typedef struct {
  pixel_t *pixels;
  size_t width;
  size_t height;
} bitmap_t;

/* The maximum iteration count. */
const int max_it = 170;

/* Compute the out-coloring based on the iteration counter. */
pixel_t outcolor(int it) {
  return 0x00010001 * ((it * 0xff) / max_it);
}

/* "Compute" the in-coloring. */
pixel_t incolor() {
  return 0x00000000;            /* black */
}

/* Computes the visual representation of a mandelbrot set for a given line. */
void compmandelbrot_line(bitmap_t* mbrot, size_t y) {

  for (size_t x = 0; x < mbrot->width; x++) {

    float complex c = ((3.0f * x / mbrot->width) - 2.0f)
        + I * ((2.0f * y / mbrot->height) - 1.0f);

    bool diverges = false;
    float complex z = 0;
    int it;
    for (it = 1; it <= max_it; it++) {
      /* z = z² + c */
      z = cpowf(z, 2) + c;

      /* If |z| ever gets greater than 2, it diverges. */
      if (cabsf(z) > 2) {
        diverges = true;
        break;
      }
    }

    uint32_t color;
    if (diverges) {
      color = outcolor(it);
    } else {
      color = incolor();
    }

    size_t i = y * mbrot->width + x;
    mbrot->pixels[i] = color;
  }
}

/* Next line to compute. */
size_t next_y = 0;

/* Mutex for next_y. */
pthread_mutex_t next_y_mutex;

/* Thread to compute mandelbrot. */
void *compmandelbrot(void *arg) {
  bitmap_t* mbrot = (bitmap_t*) arg;

  size_t y = 0;
  while (y < mbrot->height) {
    int ret;
    ret = pthread_mutex_lock(&next_y_mutex);
    assert(ret == 0);

    y = next_y++;

    ret = pthread_mutex_unlock(&next_y_mutex);
    assert(ret == 0);

    if (y < mbrot->height) {
      compmandelbrot_line(mbrot, y);
    }
  }

  pthread_exit(NULL);
}

/* Write "bitmap" to a PNG file specified by "path"; returns 0 on
   success, non-zero on error. */
static int save_png_to_file(bitmap_t * bitmap, const char *path) {
  FILE *fp;

  png_structp png_ptr = NULL;
  png_infop info_ptr = NULL;
  size_t x, y;
  png_byte **row_pointers = NULL;

  /* "status" contains the return value of this function. At first it is set to
   * a value which means 'failure'. When the routine has finished its work, it
   * is set to a value which means 'success'. */
  int status = -1;

  int pixel_size = 3;
  int depth = 8;

  fp = fopen(path, "wb");
  if (!fp) {
    goto fopen_failed;
  }

  png_ptr =
      png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
  if (png_ptr == NULL) {
    goto png_create_write_struct_failed;
  }

  info_ptr = png_create_info_struct(png_ptr);
  if (info_ptr == NULL) {
    goto png_create_info_struct_failed;
  }

  /* Set up error handling. */
  if (setjmp(png_jmpbuf(png_ptr))) {
    goto png_failure;
  }

  /* Set image attributes. */
  png_set_IHDR(png_ptr,
               info_ptr,
               bitmap->width,
               bitmap->height,
               depth,
               PNG_COLOR_TYPE_RGB,
               PNG_INTERLACE_NONE,
               PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);

  /* Initialize rows of PNG. */
  row_pointers = png_malloc(png_ptr, bitmap->height * sizeof(png_byte *));
  for (y = 0; y < bitmap->height; ++y) {
    png_byte *row =
        png_malloc(png_ptr, sizeof(uint8_t) * bitmap->width * pixel_size);
    row_pointers[y] = row;
    for (x = 0; x < bitmap->width; ++x) {
      pixel_t *pixel = bitmap->pixels + (y*bitmap->width + x);
      *row++ = (*pixel & 0x00ff0000) >> 16;
      *row++ = (*pixel & 0x0000ff00) >> 8;
      *row++ = (*pixel & 0x000000ff);
    }
  }

  /* Write the image data to "fp". */
  png_init_io(png_ptr, fp);
  png_set_rows(png_ptr, info_ptr, row_pointers);
  png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, NULL);

  /* The routine has successfully written the file, so we set "status" to a
   * value which indicates success. */
  status = 0;

  for (y = 0; y < bitmap->height; y++) {
    png_free(png_ptr, row_pointers[y]);
  }
  png_free(png_ptr, row_pointers);

png_failure:
png_create_info_struct_failed:
  png_destroy_write_struct(&png_ptr, &info_ptr);
png_create_write_struct_failed:
  fclose(fp);
fopen_failed:
  return status;
}

int main(int argc, char *argv[]) {

  /* Get command line options. */
  int num_threads = 1;
  char* options = "j:";
  if (getopt(argc, argv, options) == 'j') {
    num_threads = (int) strtol(optarg, NULL, 10);
    assert(num_threads != 0);
  } else {
    printf("Usage: multibrot -j <number_of_threads>\n");
    exit(0);
  }

  /* Set up mandelbrot bitmap. */
  bitmap_t mbrot;
  mbrot.width = 1600;
  mbrot.height = 1200;
  mbrot.pixels = calloc(sizeof(pixel_t), mbrot.width * mbrot.height);
  assert(mbrot.pixels != NULL);

  /* Start computing threads. */
  int ret;
  pthread_t* threads;
  threads = calloc(sizeof(pthread_t), num_threads);

  ret = pthread_mutex_init(&next_y_mutex, NULL);
  assert(ret == 0);

  for (int t = 0; t < num_threads; t++) {
    ret = pthread_create(&threads[t], NULL, compmandelbrot, (void *) &mbrot);
    assert(ret == 0);
  }

  /* Wait for the threads to finish. */
  for (int t = 0; t < num_threads; t++) {
    ret = pthread_join(threads[t], NULL);
    assert(ret == 0);
  }

  /* Save PNG */
  char *file = "multibrot.png";
  ret = save_png_to_file(&mbrot, file);
  assert(ret == 0);

  /* Quit. */
  pthread_exit(NULL);
  free(threads);
  free(mbrot.pixels);
  ret = pthread_mutex_destroy(&next_y_mutex);
  assert(ret == 0);
  return 0;
}