#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.
 *
 * Out-coloring means: coloring when a pixel is determined
 * to be outside of the mandelbrot set.
 *
 * @param it iteration counter when diversion was determined.
 *
 * @return pixel color
 */
pixel_t outcolor(int it) {
  return 0x00010001 * ((it * 0xff) / max_it);
}

/**
 * "Compute" the in-coloring.
 *
 * In-coloring means: coloring when a pixel is determined
 * to be inside the mandelbrot set.
 *
 * @return pixel, i.e. color
 */
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;
  }
}

/* Mandelbrot bitmap. */
bitmap_t *mbrot;

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

/* Mutex for accessing next_y. */
pthread_mutex_t next_y_mutex;

/**
 * Thread to compute mandelbrot.
 *
 * @param arg (not used, but necessary for pthread_create.)
 */
void *compmandelbrot(void *arg __attribute__((__unused__))) {
  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);
}

/**
 * Writes a bitmap to a PNG file.
 *
 * @param bitmap bitmap to save
 * @param path   PNG file to write to
 *
 * @return 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. */
  mbrot = calloc(sizeof(bitmap_t), 1);
  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 *) NULL);
    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);

  /* Clean up. */
  free(threads);
  free(mbrot->pixels);
  free(mbrot);
  ret = pthread_mutex_destroy(&next_y_mutex);
  assert(ret == 0);

  /* Quit. */
  pthread_exit(NULL);
  return 0;
}